AtCoder ABC 159 F

問題: https://atcoder.jp/contests/abc159/tasks/abc159_f

問題文が複雑ですが、\(A\)の要素のうち、\(L\)番目と\(R\)番目の間にあるものをいくつか選んで和を\(S\)にするような選び方の数が\(f(L, R)\)です。

すべての\(L, R\)について\(f(L, R)\)の和を求めるので、答えを\(Y_N\)として式にすると、

 

\[\begin{align} Y_N = \sum_{L=1}^N\sum_{R=L}^N f(L, R) \end{align}\]

 

となります。

\(N\)の部分を変数にした\(Y_i\)を変形していきます。

 

\[\begin{align} Y_i &= \sum_{L=1}^i\sum_{R=L}^i f(L, R) \\ &= \sum_{L=1}^i\sum_{R=L}^{i-1} f(L, R) + \sum_{L=1}^i f(L, i) \\ &= \sum_{L=1}^{i-1}\sum_{R=L}^{i-1} f(L, R) + \sum_{L=1}^i f(L, i) \\ &= Y_{i-1} + \sum_{L=1}^i f(L, i) \end{align}\]

 

\(g(i) = \sum_{L=1}^i f(L, i)\)とすると、\(Y_1 = \sum_{L=1}^1\sum_{R=1}^1 f(L, R) = \sum_{L=1}^1 f(L, R) = g(1)\)とあわせて、

 

\[\begin{aligned} Y_N &= Y_{N-1} + g(N) \\ &= Y_{N-2} + g(N) + g(N-1) \\ &= Y_1 + g(N) + g(N-1) + g(N-2) + ... + g(2) \\ &= \sum_{i=1}^N g(i) \end{aligned}\]

 

となるので、各\(g(i)\)を動的計画法を使って求めていきましょう。

\(A\)の要素のうち、\(L\)番目と\(R\)番目の間にあるものをいくつか選んで和を\(s\)にするような選び方の数を\(z(L, R, s)\)として（\(f(L, R) = z(L, R, S)\)です）、\(Z(i, j) = \sum_{L=1}^i z(L, i, j)\)を考えます。\(g(i) = Z(i, S)\)となるので、答えは\(Y_N = \sum_{i=1}^N Z(i, S)\)です。

 

\(Z\)には、次のような漸化式が成り立ちます。

\[Z(i, j) = \begin{cases} Z(i-1, j) &\text{if } j < A[i] \\ Z(i-1, j) + Z(i-1, j-A[i]) &\text{if } j > A[i] \\ Z(i-1, j) + i &\text{if } j = A[i] \end{cases}\]

\(i-1\)番目までの数で和を\(j\)にする選び方は、\(i\)番目までの数から選ぶときも可能です。これに加えて\(j > A[i]\)なら、和が\(j-A[i]\)になるように\(i-1\)番目までの数を選んでから、\(i\)番目の\(A[i]\)を選ぶことで和を\(j\)にできます。

\(j = A[i]\)のときは、各\(L \leq i\)に対して、\(L\)番目から\(i\)番目の数から\(i\)番目の\(A[i]\)だけを選ぶ選び方が\(i\)通りあります。

計算量は\(O(NS)\)となり、十分間に合います。

C#のソースコード：

using System;

public class Program

{

    public static void Main()

    {

        string[] input;

        input = Console.ReadLine().Split(' ');

        var n = int.Parse(input[0]);

        var s = int.Parse(input[1]);

        var a = new int[n + 1];

        input = Console.ReadLine().Split(' ');

        for (int i = 1; i <= n; i++)

        {

            a[i] = int.Parse(input[i - 1]);

        }

        var p = 998244353;

        var dp = new int[n + 1, s + 1];

        for (int i = 1; i <= n; i++)

        {

            for (int j = 1; j <= s; j++)

            {

                dp[i, j] = dp[i - 1, j];

                if (j > a[i]) dp[i, j] += dp[i - 1, j - a[i]];

                if (j == a[i]) dp[i, j] += i;

                dp[i, j] %= p;

            }

        }

        var ans = 0;

        for (int i = 1; i <= n; i++)

        {

            ans += dp[i, s];

            ans %= p;

        }

        Console.WriteLine(ans);

    }

}