【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...と探索させた。これでは途中で要素を削除した時それ以降の要素をすべて１ずつずらさなくてはいけなくなり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!!