競プロ どこが計算量を食っているのか

Question

###問題

問題リンク

（概要）
長さnの英小文字列sが与えられ、そこから最終的にちょうどk文字を取り除きたい。取り除く順番のルールは、まずaが存在する限り左から取り除く。aがなくなったらbが存在する限り左から取り除く。bがなくなったら…。yがなくなったらzが存在する限り左から取り除く。といった手法を用いる。最終的に得られるsはどのようなものか。

#コード

n, k = map(int, input().split())
s = list(str(input()))

length = 0
li = []
for i in range(97, 97 + 26):
	for j in range(n):
		if s[j] == chr(i):
			li.append(j)
			length += 1
		if length == k:
			break
	else:
		continue
	break

li = set(li)

ans = ''
for i in range(n):
	if i not in li:
		ans += s[i]

print(ans)

#方針
aからzの順に、かつ左から右の順に、取り除かれる要素のインデックスをリストliに追加していき、最後はそのli（setに変換後）にインデックスが含まれていないものだけを文字列ansに追加していき出力する。

#疑問
制限時間は2秒となっており、10＾8回程度の計算は耐えられるかなと思い、nが高々4×10^5であることに鑑みて、a-zとの二重ループでも問題はないだろうという風に鷹を括っていたのですが、どの部分が具体的に計算時間を増幅させてしまっているのでしょうか。それとも、そもそもの話としてもう少し余裕をもったコードに書き換えた方が良いのでしょうか（その場合は一から考え直すのですが…）。素人質問にて恐縮ですが、お力添えいただける点がございましたら、何卒よろしくお願い申し上げます。

Answer 1

otnさんの言う通りで、プロファイラを使いましょう。

適当に40万文字から20万文字を削除するデータを作って、

$ python3 -m cProfile -s tottime test.py < test.txt

のように実行したら、このように表示されました。

         5408377 function calls in 8.425 seconds

   Ordered by: internal time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    4.754    4.754    8.425    8.425 test.py:1(<module>)
  5208223    3.515    0.000    3.515    0.000 {built-in method builtins.chr}
   200000    0.135    0.000    0.135    0.000 {method 'append' of 'list' objects}
        2    0.015    0.007    0.015    0.008 {built-in method builtins.input}
        1    0.005    0.005    0.005    0.005 {built-in method builtins.print}
       49    0.000    0.000    0.000    0.000 {built-in method _codecs.utf_8_decode}
       49    0.000    0.000    0.001    0.000 codecs.py:319(decode)
       49    0.000    0.000    0.000    0.000 codecs.py:331(getstate)
        1    0.000    0.000    8.425    8.425 {built-in method builtins.exec}
        1    0.000    0.000    0.000    0.000 {method 'split' of 'str' objects}
        1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}

これを見れば、chrがあやしいと分かるでしょう。

追記

こちらの環境だと、chrの修正だけで1秒強にまで縮んだのですが、不十分でしたか。
もう少しだけ高速化してみましたが、ここまでやらなくてはいけない時点で何かおかしい気はします。

• リストを事前確保してappendをやめる (プロファイラを外したらほとんど効果がなかったので取り下げます)
• if length == k:の判定を、lengthが増えた時だけにする

n, k = map(int, input().split())
s = list(str(input()))

length = 0
li = []
for i in range(97, 97 + 26):
    c = chr(i)
    for j in range(n):
        if s[j] == c:
            li.append(j)
            length += 1
            if length == k:
                break
    else:
        continue
    break

li = set(li)

ans = ''
for i in range(n):
    if i not in li:
        ans += s[i]

print(ans)

Comments

[kay_ventris4]

後ほどプロファイラについては理解を深めようとしておりましたが、その見方まで示して頂きありがとうございました。

built-in関数を呼び出す回数を減らす手法で、再度トライして見たいと思います。

[otn]

chrはループの外に出せますが、appendは無理ですね。

[kay_ventris4]

今、ループ中でのchr()の呼び出しを避けるため、あらかじめ
dic = {}
for i in range(97, 97 + 26):
　　dic[i] = chr(i)
としてアルファベット参照用の辞書を用意し、ループ中のchr(i)をdic[i]に変えたのですが、問題を提出したところ同じく2000msオーバーとなってしまいました。この方法ではchr()の呼び出し回数を減らしたことにならないのでしょうか。

[otn]

ループの外に出すというのは、
for i in range(97, 97 + 26):
＿ci = chr(i)
＿for j in range(n):
＿＿if s[j] == ci:
です。

[kay_ventris4]

残念ながら、こちらに変えても時間切れとなりました。
あとは疑わしいとしたらやはり.append()くらいしかないでしょうか…

[kay_ventris4]

一応自己解決の欄に別途累積和を用いたコードを貼ったのですが、こちらが正解扱いとなったのは、.append()による処理回数が非常に少ないという特徴があります。

Answer 2

どの部分が具体的に計算時間を増幅させてしまっているのでしょうか。

こういうツールを使います。
Python プロファイラ — Python 3.10.0b2 ドキュメント

#追記
違う考え方でのコード。

n, k = map(int, input().split())
s = list(str(input()))

a = [0]*(ord("z")+1)
for i in s:
    a[ord(i)] += 1

b = 0
for c in range(ord("a"),ord("z")+1):
    b += a[c]
    if b >= k:
        break
d = a[c]-(b-k)

ans = ""
for i in s:
    e = ord(i)
    if e > c:
        ans += i
    elif e == c:
        if d > 0:
            d -= 1
        else:
            ans += i
print(ans)

Comments

[otn]

と、回答してから質問文のコードを見ましたが、この短さ・単純さであれば調べるまでも無いかと。

[kay_ventris4]

ご回答ありがとうございます。
プロファイラーというものについては初見でしたので、後ほどゆっくりと時間をかけて勉強させて頂きたいと思います。

質問文のコードについてですが、このままの提出だと時間切れを起こしてしまうのは、どのような原因が考えられるか、ご教授頂くことは可能でしょうか。

[otn]

2秒で終わらないデータというのは n と k はどれくらい？

[kay_ventris4]

n, k = 400000, 123456 のテストケースで実行時間が2000msを超えてしまいました。

[otn]

n が大きいとすると、s を1文字ずつスキャンする回数を抑えないといけないので、
2回だけにする方法でやってみると、
（最初のスキャンで削除する文字範囲を決めた後、2回目で実際に削除処理）
n, k = 218525 123456
のデータで、
質問のコードで
real 0m0.732s
user 0m0.719s
sys 0m0.009s
が、こうなりました。
real 0m0.127s
user 0m0.121s
sys 0m0.005s

[kay_ventris4]

追記頂いたコードを試しに提出してみたのですが、実行時間超過となりました。

[otn]

sを最後まで見ないと削除する文字範囲が決まらないため、2回スキャンするのは必須だと思うので、あとは細かいチューニングでしょうか。
あらかじめ、
t = [ord(x) for x in s]
とordしたリストを作ってそれを使うことで数%は速くなりますね。

[otn]

回答に吊られて、range使ったけど、for i in s: にすれば速いか。
書き直しておきます。

[kay_ventris4]

ありがとうございます。

Answer 3

ヒントを欲しいという質問ではないようですので、Yes/Noで回答します。

• そもそもの話としてもう少し余裕をもったコードに書き換えた方が良いのでしょうか

はい。

Comments

[kay_ventris4]

ご回答ありがとうございます。
逆に言えば、該当コード中のここさえ直せば計算時間を大きく減らせるという箇所は特にないという理解でよろしいでしょうか。

Answer 4

一応、累積和などを使ってぐちゃぐちゃ書いていたら正解はしました。
しかし、上のコードの時間切れ理由は未だわからずという状況です。

from itertools import accumulate

n, k = map(int, input().split())
s = list(str(input()))

dic = {}
for i in range(97, 97 + 26):
	dic[chr(i)] = 0
for el in s:
	dic[el] += 1

li = list(dic.values())
li_ = list(accumulate(li))

final = float('inf')
for i in range(26):
	if li_[i] > k:
		final = i
		break

if final != 0 and final != float('inf'):
	remove = []
	for i in range(final):
		remove.append(i)
	remove = set(remove)

	s = [el for el in s if ord(el) - 97 not in remove]

	final_length = k - li_[final - 1]
	final_remove = []
	length = 0
	for i in range(len(s)):
		if length == final_length:
			break
		if s[i] == chr(final + 97):
			final_remove.append(i)
			length += 1
	final_remove = set(final_remove)

	s = [s[i] for i in range(len(s)) if i not in final_remove]

	print(''.join(s))

elif final == 0:
	bye = []
	for i in range(len(s)):
		if s[i] == 'a':
			bye.append(i)
	bye = bye[:k]
	bye = set(bye)

	s = [s[i] for i in range(len(s)) if i not in bye]

	print(''.join(s))

elif final == float('inf'):
	print('')