AtCoder ARC059 E - キャンディーとN人の子供 / Children and Candies (Go)

irisruneです。今回は部分点込みの解説になるのでコードの分量が多めなのと、式を使った解説が多くなっています。

問題

部分点解法は典型的な…？DP問題、満点解法はその発展となっています。

部分点解法（400点）

package main

import (
    "bufio"
    "fmt"
    "os"
    "strconv"
)

var sc *bufio.Scanner

func nextInt() int {
    sc.Scan()
    i, e := strconv.Atoi(sc.Text())
    if e != nil {
        panic(e)
    }
    return i
}

const MOD = 1000000007

func main() {
    sc = bufio.NewScanner(os.Stdin)
    sc.Split(bufio.ScanWords)
    n, c := nextInt(), nextInt()

    aArray := make([]int, n)
    for ni := range aArray {
        aArray[ni] = nextInt()
    }
    bArray := make([]int, n)
    for ni := range bArray {
        bArray[ni] = nextInt()
        if aArray[ni] != bArray[ni] {
            fmt.Println(-1)
            return
        }
    }

    dp := make([][]int, c+1)
    for ci := range dp {
        dpLine := make([]int, n)
        dp[ci] = dpLine
    }

    for ci, dpLine := range dp {
        for ni := range dpLine {
            if ci == 0 {
                dp[ci][ni] = 1
                continue
            }
            if ni == 0 {
                dp[ci][ni] = dp[ci-1][ni] * aArray[ni] % MOD
                continue
            }
            dp[ci][ni] = (dp[ci][ni-1] + (dp[ci-1][ni] * aArray[ni] % MOD)) % MOD
        }
    }
    fmt.Println(dp[c][n-1])
}

コードを見ればわかる通り $A_i=B_i$ の場合は単純なDPで解くことができます。計算量は $O(CN)$ となるため十分間に合います。

DPで解けることを示す前に全探索できないか考えてみましょう。区別できない $C$ 個のキャンディーを $N$ 人の子供に配る場合の数は $_{N+C-1}C_C$ 通りです¹。これは $N=400,C=400$ のとき非常に大きな数となるため、全探索は現実的ではありません。

そこでDPで解くことを考えます。 $C=C_0,N=N_0$ のときの $f(x_1,x_2,\ldots,x_N)$ の値を $F(C_0)(N_0)$ とおきます。このとき、 $F(C_0)(N_0)=x_{N_0}F(C_0-1)(N_0)+F(C_0)(N_0-1)$ が成り立つのでDPにより解くことができます。

満点解法

package main

import (
    "bufio"
    "fmt"
    "os"
    "strconv"
)

var sc *bufio.Scanner

func nextInt() int {
    sc.Scan()
    i, e := strconv.Atoi(sc.Text())
    if e != nil {
        panic(e)
    }
    return i
}

const MOD = 1000000007
const UBOUND = 400

func main() {
    sc = bufio.NewScanner(os.Stdin)
    sc.Split(bufio.ScanWords)
    n, c := nextInt(), nextInt()

    aArray := make([]int, n)
    for ni := range aArray {
        aArray[ni] = nextInt()
    }
    bArray := make([]int, n)
    for ni := range bArray {
        bArray[ni] = nextInt()
    }

    powMat := make([][]int, UBOUND+1)
    for r := range powMat {
        pow := make([]int, UBOUND+1)
        for i := range pow {
            if r == 0 {
                pow[i] = 1
                continue
            }
            pow[i] = (powMat[r-1][i] * i) % MOD
        }
        powMat[r] = pow
    }
    powCumSumMat := make([][]int, UBOUND+1)
    for r := range powCumSumMat {
        powCumSum := make([]int, UBOUND+1)
        for x := range powCumSum {
            if x == 0 {
                powCumSum[x] = powMat[r][x]
                continue
            }
            powCumSum[x] = (powCumSum[x-1] + powMat[r][x]) % MOD
        }
        powCumSumMat[r] = powCumSum
    }

    dp := make([][]int, c+1)
    for ci := range dp {
        dpLine := make([]int, n)
        dp[ci] = dpLine
    }

    for ci, dpLine := range dp {
        for ni := range dpLine {
            if ni == 0 {
                dp[ci][ni] = powCumSumMat[ci][bArray[ni]] - powCumSumMat[ci][aArray[ni]-1]
                switch {
                case dp[ci][ni] < 0:
                    dp[ci][ni] += MOD
                default:
                    dp[ci][ni] %= MOD
                }
                continue
            }
            for cj := 0; cj <= ci; cj++ {
                switch {
                case powCumSumMat[cj][bArray[ni]]-powCumSumMat[cj][aArray[ni]-1] < 0:
                    dp[ci][ni] += dp[ci-cj][ni-1] * (powCumSumMat[cj][bArray[ni]] - powCumSumMat[cj][aArray[ni]-1] + MOD) % MOD
                default:
                    dp[ci][ni] += dp[ci-cj][ni-1] * ((powCumSumMat[cj][bArray[ni]] - powCumSumMat[cj][aArray[ni]-1]) % MOD) % MOD
                }
                dp[ci][ni] %= MOD
            }
        }
    }
    fmt.Println(dp[c][n-1])
}

こちらもDPをベースとした解法となります。計算量は $O(C\max(A_i,B_i)+C^2N)$ となりますが、制限時間が4秒なのでまだ間に合う範囲です。

まずは $f(x_1+x_2+\ldots+x_N)$ について部分点解法と同様に求める方針を考えてみましょう。その場合全体の計算量は $O(CN\times (\max(A_i,B_i))^N)$ となってしまい到底間に合いません。そのため $\sum^{B_1}_{A_1}\ldots\sum^{B_N}_{A_N}f(x_1+x_2+\ldots+x_N)$ を一つの塊としてDPで求める必要があります。

部分点解法と同様に、 $C=C_0,N=N_0$ のときの $\sum^{B_1}_{A_1}\ldots\sum^{B_N}_{A_N}f(x_1+x_2+\ldots+x_N)$ の値を $F(C_0)(N_0)$ とおきます。すると、 $F(C_0)(N_0)=(F(C_0-0)(N_0-1))(A_{N_0}^0+(A_{N_0}+1)^0+\ldots+B_{N_0}^0)$ $+(F(C_0-1)(N_0-1))(A_{N_0}^1+\ldots+B_{N_0}^1)+\ldots$ $+(F(C_0-C_0)(N_0-1))(A_{N_0}^{C_0}+\ldots+B_{N_0}^{C_0})$ が成り立ちます。²

あとはこの関係式を基に実装をすればよいですが、注意点が2つあります。

愚直に実装すると計算量が $O(C^2N\times C\max(A_i, B_i))$ となってしまうので、累積和を前計算で求める（計算量 $C\max(A_i, B_i)$ ）ことによりDPの計算量を $O(C^2N)$ に抑える必要があります。
累積和から部分和を計算する際に減算を行うため、剰余の計算のために正負の場合分けを行う必要があります。

以上2点に注意すればDPにより解を求めることができます。

雑記

そもそもDPに帰着できないと部分点を取ることすらままならないですが、くらいの簡単なケースで式をこねくり回すような道筋しか思いつきません。
- 同じように式をこねくり回したら満点解法まで辿り着けたので、アプローチとしては間違ってないのかもしれませんね。

具体的な計算方法は長くなるので重複組合せで検索していただければと思います。↩
$x_{N_0}=A_{N_0}=B_{N_0}$ で部分点の条件を満たすとき、 $F(C_0)(N_0)=x_{N_0}^0F(C_0-0)(N_0-1)+(x_{N_0}^1F(C_0-1)(N_0-1)+\ldots))$ $=F(C_0)(N_0-1)+x_{N_0}F(C_0-1)(N_0)$ と変換できるので、部分点解法と同様の関係式となります。↩