【ABC380】Pythonで文字列結合をする際はjoinを使うこと

問題C - Move Segment

https://atcoder.jp/contests/abc380/tasks/abc380_c

こちらの問題で、アルゴリズム的には適切な時間で実行されるのにも関わらず、PythonだとTLEが出てしまった。

問題のコード

そのときのコードが下記の通り。
アルゴリズムは「解説」の「コンテスト全体の解説」の動画参照。

N, K = map(int, input().split())
S = input()

arr = []
for i in range(N):
    ch = S[i]
    if i > 0 and S[i] == S[i - 1]:
        arr[-1][-1] += 1
    else:
        arr.append([ch, 1])

one_cnt = 0
for i in range(len(arr)):
    ch, cnt = arr[i]
    if ch == "1":
        one_cnt += 1
    if one_cnt == K:
        arr[i], arr[i - 1] = arr[i - 1], arr[i]
        break

ans = ""
for ch, cnt in arr:
    ans += ch * cnt
print(ans)

実行時間は $O(N)$ で、$1 \le N \le 10^5$ なので問題ないはず。にも関わらず実際に実行すると「2227ms」かかった。

修正コード

変えたのはans宣言以降。

N, K = map(int, input().split())
S = input()

arr = []
for i in range(N):
    ch = S[i]
    if i > 0 and S[i] == S[i - 1]:
        arr[-1][-1] += 1
    else:
        arr.append([ch, 1])

one_cnt = 0
for i in range(len(arr)):
    ch, cnt = arr[i]
    if ch == "1":
        one_cnt += 1
    if one_cnt == K:
        arr[i], arr[i - 1] = arr[i - 1], arr[i]
        break

# ここから変更
ans = []
for ch, cnt in arr:
    ans.append(ch * cnt)
print("".join(ans))

文字列結合を += から join に変えただけである。これによって実行時間はおよそ1/10である「200ms」になった。

原因

Pythonで文字列はイミュータブルなので、+= を使うたびに新しいオブジェクトが作成される。ここのコストが余分にかかってしまう。そのためjoinよりも時間がかかってしまう。