【Python】AtCoderでリストにpop()したところTLEが発生したので…
初めに
今回やっていたのはAtCoder_ABC_461でA,Bと解き進めていたのだがCで手が都止まり最終的に完成したコードがTLEになったのでそのコードの振り返りで気づいたことを書く。
修正コード
まず自分が作ったコードは、
# TLEになってしまう危険な書き方
while check < m:
t = j_list[-i]
if need_[t[1]] == 1:
i += 1
continue
else:
ans += t[0]
need_[t[1]] = 1
# ここが問題!!
j_list.pop(-i)
check += 1
で、問題の一行はj_list.pop(-i)である。今回は後ろから一個ずつ確認し条件を満たすまで反復するコードとしてpop()の引数に-1,-2,-3...と探索させた。これでは途中で要素を削除した時それ以降の要素をすべて1ずつずらさなくてはいけなくなりO(n)かかる。
改善案としては消すのではなく使用済みを示すマークにすることが考えられた。
考えたコードは
# 改善版:フラグを使った安全な書き方
# 要素数と同じサイズの「使用済みメモ」を用意
used_item = [False] * n
# 1. まず M 種類の色の最大価値を確保する
for i in range(n):
v = j_list[i][0]
c = j_list[i][1]
if not need_[c]:
ans += v
need_[c] = True
used_item[i] = True # 「この宝石は使ったよ」とメモ!
check += 1
if check == m:
break
# 2. 残りの枠を、まだ使っていないものから取る
count = 0
for i in range(n):
if count == k - m:
break
if not used_item[i]: # メモを見て、まだ使っていないなら取る
ans += j_list[i][0]
count += 1
で、Trueに置き換えていくことで瞬時に計算できる
まとめ
自分はpop()について、探索するときにインデックスで一発で探索できるかO(1)と考えていたが、それ以降のlistを整理する際の処理でO(n)かかることを知った。
これからは新しい知識や今後したいことなど、書いていこうと思う。目指せ緑AtCoderer!!