[学習記録]AtCoder EDP M-Candies[Go]

問題

N 人の子供たちがいます。 子供たちには 1,2,…,N と番号が振られています。

子供たちは K 個の飴を分け合うことにしました。 このとき、各 i (1≤i≤N) について、子供 i が受け取る飴の個数は 0 以上 ai​ 以下でなければなりません。 また、飴が余ってはいけません。

子供たちが飴を分け合う方法は何通りでしょうか？ 10^9+7 で割った余りを求めてください。 ただし、2 通りの方法が異なるとは、ある子供が存在し、その子供が受け取る飴の個数が異なることを言います。

解法

i人の子供で、j 個の飴を分け合うことを考える。
j　個の飴を分け合う組み合わせは i 番目の子供がとる数によって変わっていき、その数は i -1番目の子供たちが取った飴の数を、i 番目の子供が取る数を変えていったときの総和になる

// i=2, j=0,  (0,0)
// i=2, j=1,  (0,1)(1,0)
// i=2, j=2,  (1,1)(0,2)
// i=2, j=3,  (1,2)
// i=2, j=4,

// i=3, j=0,  (0,0,0)
// i=3, j=1,  (1,0,0)(0,1,0) | (0,0,1)
//  > i=3の子供が0ことるとき、組み合わせは i=2, j=1と同じになる
//  > (0,1,[0]),(1,0,[0])
//  > i=3の子供が1ことるとき、組み合わせは i=2, j=0と同じになる
//  > (0,0,[1])
//  >  2 + 1 = 3
// i=3, j=2,  (1,1,0) | (1,0,1)(0,1,1)
// i=3, j=3,  (1,2,0) | (1,1,1)(0,2,1) | (1,0,2)(0,1,2) | (0,0,3)
// i=3, j=4,  (1,2,1) | (0,2,2)(1,1,2) | (0,1,3)(1,0,3) |

dp := make([][]int, n+1)
dp[0] = make([]int, k+1)
dp[0][0] = 1

for i := 1; i <= n; i++ {
	dp[i] = make([]int, k+1)

	for j := 0; j <= k; j++ {
        // i 番目の子供が取る数を変えていったときの総和
		dp[i][j] = Σdp[i-1][x]
	}
}

ans := dp[n][k]

return ans

しかしながら、Σdp[i-1][x]をそのまま計算しようとすると、計算量が
O（NK^2)となり、TLEしてしまう。
そのため、累積和を取り、計算量を減らす。（O(NK))

sli := readSli(n)

for i := 1; i <= n; i++ {
	dp[i] = make([]int, k+1)
	cum := make([]int, k+2)

	cum[0] = 0
	for j := 1; j <= k+1; j++ {
		cum[j] = (cum[j-1] + dp[i-1][j-1])
	}

	for j := 0; j <= k; j++ {
		dp[i][j] = Σdp[i-1][x]
	}
}

Σの部分は、i 番目の人が何個受け取れるかに依存するのでそれに依存する
受け取れる個数( a ) が k より少ない場合、i 番目の人が受け取る個数は
k -a < j < k の間で変化し、これを累積和を用いて表すと
dp[i][j]   =    Σdp[i-1][x]　＝　cum[ j + 1] - cum[ j - a ]         になる。
( j + 1 ・・・累積和のため、j まで足したときの和　= j＋１となる）
累積和の開始インデックスがゼロを下回ることは無いため、Maxを取る

for i := 1; i <= n; i++ {
	dp[i] = make([]int, k+1)
	cum := make([]int, k+2)

	cum[0] = 0
	for j := 1; j <= k+1; j++ {
		cum[j] = (cum[j-1] + dp[i-1][j-1])
	}

	for j := 0; j <= k; j++ {
		dp[i][j] = cum[j+1] - cum[Max(0, j-sli[i-1])]
	}
}

dpのn, k が回答になる

コード（AC)

func run() interface{} {
	n, k := readInt(), readInt()

	sli := readSli(n)
	dp := make([][]int, n+1)
	dp[0] = make([]int, k+1)
	dp[0][0] = 1

	for i := 1; i <= n; i++ {
		dp[i] = make([]int, k+1)
		cum := make([]int, k+2)

		cum[0] = 0
		for j := 1; j <= k+1; j++ {
			cum[j] = (cum[j-1] + dp[i-1][j-1]) % prime_number
		}

		for j := 0; j <= k; j++ {
			dp[i][j] = (cum[j+1] - cum[Max(0, j-sli[i-1])] + prime_number) % prime_number
		}
	}

	ans := dp[n][k]

	return ans
}