実行時間超過 原因は定数倍？

Question

#問題

AtCoder Regular Contest 104 B問題より

#方針
1<=N<=5000という制約から、Sの部分文字列の全列挙自体は計算量は5000*(5000 + 1)//2と見積もられ、実行制限時間2秒を超えることはないと想定されます。しかし、実際に部分文字列を取り出そうとするとスライスの部分で余計な計算量を食うので、あらかじめi番目までにA, T, G, Cが何回登場したかを累積和を用いてA_rec, T_rec, G_rec, C_recに格納しておき、S[i:j]なる部分文字列についてそこから必要な値を用い、その部分文字列に何が何個含まれているかを調べます。そして(Aの個数)=(Tの個数)かつ(Gの個数)=(Cの個数)となれば良いので、それぞれ調べカウントします。

#TLEコード

from itertools import accumulate

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

A_rec = [0] * N
T_rec = [0] * N
G_rec = [0] * N
C_rec = [0] * N

for i in range(N):
	if S[i] == 'A':
		A_rec[i] += 1
	elif S[i] == 'T':
		T_rec[i] += 1
	elif S[i] == 'G':
		G_rec[i] += 1
	elif S[i] == 'C':
		C_rec[i] += 1
A_rec = list(accumulate(A_rec))
T_rec = list(accumulate(T_rec))
G_rec = list(accumulate(G_rec))
C_rec = list(accumulate(C_rec))

def substring(S):
    li = []
    for i in range(N):
        for j in range(i + 1, N + 1):
        	if i > 0:
        		li.append([A_rec[j - 1] - A_rec[i - 1], T_rec[j - 1] - T_rec[i - 1], G_rec[j - 1] - G_rec[i - 1], C_rec[j - 1] - C_rec[i - 1]])
        	elif i == 0:
        		li.append([A_rec[j - 1], T_rec[j - 1], G_rec[j - 1], C_rec[j - 1]])
    return li

cnt = 0
for el in substring(S):
	if el[0] == el[1] and el[2] == el[3]:
		cnt += 1

print(cnt)

#質問
substringのところで二重ループを用いましたが、これはNの制約から直接的に実行時間超過に起因しているわけではなさそうであると推察しています。またそれ以外の処理については、全て線形時間で行っているので、これもまた大きな計算量を食っているようには思えません。実行はPypy3とPython3の両方で試しましたが、やはりPythonという言語の定数倍計算の処理の遅さがTLEの原因なのでしょうか。それとも、何か見落としている大きな計算がありますでしょうか。素人質問にて恐縮ですが、お力添え頂ける箇所がございましたら、ご教授のほどよろしくお願い申し上げます。

Comments

[BeatStar]

画像で質問するより、URLでの質問のほうがいいと思いますよ。

Answer 1

私の見る限り、問題は二つあります。

問題１　アルゴリズム
不要な計算が多いと思います。
例：　長さが奇数の時に必要もないのに調べている。
他もいろいろあってアルゴリズムの根本的な見直しをすべきです。
それを角と、この問題のネタバレになるので、まずはご自身で考えてください。

問題２ メモリ割当ての多いこと
S[i]を使うと、新しいstrのインスタンスが作られるので結構時間を使います。
li.appendを使うと、メモリのとりなおしが数回に一回おきるので結構時間を使います。

これらの二つの問題をなくせばかなりスピードアップするでしょう。

Comments

[kay_ventris4]

自己解決の欄にコードは載せましたが、ご教示頂いた通りj-iが奇数になる場合を飛ばし、さらにappendの過程を抜いて直接関数内でカウントすることにより実行時間が最大183msまで削ることができました。後ほどアルゴリズムの自己評価は行いますが、とりあえずこれだけでも大きな成果でした。ありがとうございました。

Answer 2

cProfile
https://docs.python.org/ja/3/library/profile.html
とかを使うとどこが遅いかだいたいわかると思いますよ。

この使い方が便利なのではないかと思います。

python -m cProfile [-o output_file] [-s sort_order] (-m module | myscript.py)

ご参考まで。

Answer 3

リストへのappendを省き、(j - i)が偶数である時のみのアルゴリズムに変更したところ、うまく行きました：

from itertools import accumulate

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

A_rec = [0] * N
T_rec = [0] * N
G_rec = [0] * N
C_rec = [0] * N

for i in range(N):
	if S[i] == 'A':
		A_rec[i] += 1
	elif S[i] == 'T':
		T_rec[i] += 1
	elif S[i] == 'G':
		G_rec[i] += 1
	elif S[i] == 'C':
		C_rec[i] += 1
A_rec = list(accumulate(A_rec))
T_rec = list(accumulate(T_rec))
G_rec = list(accumulate(G_rec))
C_rec = list(accumulate(C_rec))

def solve(S):
	cnt = 0
	for i in range(N):
		for j in range(i + 1, N + 1):
			if (j - i) % 2 == 0:
				if i > 0:
					if A_rec[j - 1] - A_rec[i - 1] == T_rec[j - 1] - T_rec[i - 1] and G_rec[j - 1] - G_rec[i - 1] == C_rec[j - 1] - C_rec[i - 1]:
						cnt += 1
				elif i == 0:
					if A_rec[j - 1] == T_rec[j - 1] and G_rec[j - 1] == C_rec[j - 1]:
						cnt += 1
	return cnt

print(solve(S))