AtCoder ABC158 高速化の方法と計算量の見積もり

Question

解決したいこと

以下のコードを提出しましたが、TLEとなってしまいました。
(44ファイル中12ファイルTLE)
計算量を改善するにはどうすべきだったのでしょうか。

問題は以下の通りです。
https://atcoder.jp/contests/abc158/tasks/abc158_e

疑問点と知りたいこと

自分の予想では、Numpyを使えば繰り返し内部の計算量はO(1)程度だと予測し、全体の計算量はO(N)程度であり、問題ないと考えていました。
しかし、実際にはTLEとなってしまい、O(10^9)以上の計算量が必要となっていたと考えられます。
このコードの計算量はどれほどだったのでしょうか。また、それはどのようにして予測できたのでしょうか。

該当のソースコード

s = list(map(int, input().strip()))

dp = np.zeros(p,dtype = int) #動的計画法

count = 0
d = 1
for i in range(n): #1の位から順に数を追加していく
    ne_ar = np.zeros(p,dtype = int)
    if(p == 2 or p == 5): #10と互いに素ではない場合
        ne = int(s[n-i-1]) % p
        if(ne == 0):
            dp[0] += 1
        count += dp[0] #
    else:
        ne = (int(s[n-i-1]) * d) % p #10とは互いに素なので、sと同じ位で考えて良い
        ne_ar[ne:] += dp[:(p-ne)] #新しく追加する数の余りをずらす
        ne_ar[:ne] += dp[(p-ne):]
        ne_ar[ne] += 1 #新しい数のみの時
        dp = ne_ar
        d *= 10
        d %= p
        count += ne_ar[0] #neを一番上の位とする部分集合の数を足す

print(count)

Answer 1

        ne_ar[ne:] += dp[:(p-ne)] #新しく追加する数の余りをずらす
        ne_ar[:ne] += dp[(p-ne):]
        ne_ar[ne] += 1 #新しい数のみの時

配列をコピーしているので、ここで時間計算量 O(P) 掛かりそうな気がします。
numpyは詳しくないので実は参照が使われて計算量O(1)になっていたらごめんなさい

Comments

[shake9]

自分もNumpyには詳しくないので分からないです。
ただ、以下の例ではNumpyを使うことで単純な加算を行う場合に対して大幅な時間短縮ができました。
そのため、O(P)よりは小さいのでは、と思っています。(O(1)なのかどうかは分かりません)
https://teratail.com/questions/266785

[maai]

しまった、+= はベクトル加算なのでコピーでは無いですね…
266785 は、 Pythonのリストをnumpy arrayに置き換えることで数十倍高速化する感じなので、入力に対する計算量オーダが改善している訳では無いです。
N要素の単純なベクトル加算は、どう頑張ってもNに比例する時間が掛かってしまいます。

もう1箇所気づきました。
for i in range(n): #1の位から順に数を追加していく
ne_ar = np.zeros(p,dtype = int)

np.zeros(p,dtype = int)で0埋めされた長さpの配列を確保しています。
これの時間計算量がO(P)、この行が呼ばれる回数がN回なので、時間計算量はO(NP)以上となり、制限時間に間に合わない気がします。

[shake9]

ありがとうございます。
そうなのですね、納得できました。
この方法ではどう頑張っても長さPの配列をN回用意する必要がありそうなので、別の方法を検討することにします。