問題
https://atcoder.jp/contests/dp/tasks/dp_m
解説を見ると、累積和で解く、ということで、計算量はO(NK)となっていますが、
この計算量ではpythonを使って解くと入力例4ですらTLEになってしまいます。
下記が累積和を使って書いたコードです。
python
1import numpy 2 3#入力の取得 4f1 = open(r"C:\tmp\input4.txt") 5f2,f3 = f1.readlines() 6n = int(f2.split()[0]) 7k = int(f2.split()[1]) 8a = list(map(int,f3.split())) 9mod = 10**9 + 7 10 11#dpは2次元配列 12#dp[i][j]→i人目まで(0-Indexed)の子供たちにj個の飴を配る時の、場合の数 13dp = [[0]*(k+1) for i in range(n)] 14 15#初期値 16#1人の子供にx個の飴を配る場合の数→それぞれ1通り 17for i in range(a[0]+1): 18 dp[0][i] = 1 19 20for i in range(1,n): 21 for j in range(k+1): 22 #累積和、最初は0 23 s = numpy.cumsum(dp[i-1]).tolist() 24 s.insert(0, 0) 25 dp[i][j] = (s[j+1] - s[max(0,j-a[i])])%mod 26print(dp[n-1][k]) 27
上のコードは計算量O(NK)であってますか?
また、他の方のコードを拝見しますと、
python
1mod = 10**9+7 2import numpy as np 3f1 = open(r"C:\tmp\input1.txt") 4f2,f3 = f1.readlines() 5n, k = map(int,f2.split()) 6A = tuple(map(int, f3.split())) 7dp = np.zeros(k+1, np.int64) 8dp[0] = 1 9for a in A: 10 dp[a+1:] -= dp[:-(a+1)] 11 dp = dp.cumsum() 12 dp %= mod 13print(dp[k])
となっており、入力例1で処理を追ってみると、
初期値dp=[1,0,0,0,0]と、子供iに配るmaxの個数の配列A=[1,2,3]があって、
dpをA[0]分ずらして引き算して、累積和を取る、ということをしています。
だからdpは下記の通り推移しまして、
dp=[1,0,0,0,0]
↓
dp=[1,1,0,0,0]
↓
dp=[1,2,2,1,0]
↓
dp=[1,3,5,6,5]
最終的なdpは解答に必要な、私のコードで言うdp[2]の列に等しくなっています。
このコードだとpythonでも時間切れにならずに通りますが、
このコードの計算量はO(NK)ですか?
このコードに至る考え方をご教示ください。
バッドをするには、ログインかつ
こちらの条件を満たす必要があります。