ABC380 C - Move segments で、TLEが発生します。

実現したいこと

ABC380 C - Move segments(https://atcoder.jp/contests/abc380/tasks/abc380_c)
で、0が連続している部分と1が連続している部分をSblockによってリスト化し、問題文通りにswapするアルゴリズムを組みました。

発生している問題・分からないこと

計算量はo(10^6)だと思うのですが、なぜか2つのケースでのみTLEします。
みんなの解答や公式解答を見る限り、特別悪い実装ではないと思うのですが、いかがでしょう？

該当のソースコード

python
1def list_int(numbers):
2    numbers_list = map(int, numbers.split(" "))
3    # 数を入力する スペースで区切る　int型にしている
4    numbers_list = list(numbers_list)
5    return numbers_list
6
7
8N, K = list_int(input())
9
10S = input()
11
12Sblock = []
13block = ""
14
15for i in range(len(S)):
16    if block == "":
17        block = S[i]
18        continue
19    if block[0] == "0" == S[i] or block[0] == "1" == S[i]:
20        block = "".join([block, S[i]])
21    else:
22        Sblock.append(block)
23        block = S[i]
24
25if block != "":
26    Sblock.append(block)
27    block = ""
28
29
30onecnt = 0
31for i in range(len(Sblock)):
32    if Sblock[i][0] == "1":
33        onecnt += 1
34    if onecnt == K:
35        swap1 = Sblock[i - 1]
36        swap2 = Sblock[i]
37        Sblock[i - 1] = swap2
38        Sblock[i] = swap1
39        break
40
41S_ans_str = "".join(Sblock)
42print(S_ans_str)
43

試したこと・調べたこと

teratailやGoogle等で検索した
ソースコードを自分なりに変更した
知人に聞いた
その他

上記の詳細・結果

文字列連結のメソッドを＋からjoinに変更しましたが、うまくいきませんでした。

補足

特になし

行動規範の内容に同意します

回答1件

ベストアンサー

blockに文字列を1文字ずつ足し合わせているのが、遅い原因です。

Pythonで、文字列を1文字ずつ足し合わせて作成すると、文字列の長さをNとしてO(N^2)の時間がかかります。(ただし、こちらの記事によると、条件によってはO(N)となる場合もあるそうです)

例えば、ご提示のコードで、8-10行目を以下のように書き換えて実行すると、時間がかかるのが分かると思います。

python
1N, K = 500000, 2
2S = "10" + "1" * 499998

公式回答では文字列をスライスで取り出していますし、他の人の回答でも1文字ずつ足し合わせるような実装は避けているはずです。

blockを文字列のリストにして、最後にjoinで結合するようにすれば、高速化します(公式のFAQも参照)。12-27行目を以下のように書き換えて実行してみてください。

python
1Sblock = []
2block = []
3
4for i in range(len(S)):
5    if block == []:
6        block = [S[i]]
7        continue
8    if block[0] == "0" == S[i] or block[0] == "1" == S[i]:
9        block.append(S[i])
10    else:
11        Sblock.append("".join(block))
12        block = [S[i]]
13
14if block != []:
15    Sblock.append("".join(block))
16    block = []