suta精進記

競プロ関連体験記や精進記録などを書きます

AGC049 D : Convex Sequence

問題リンク

atcoder.jp

問題概要

整数$ N, M $が与えられる。長さが$N$かつ以下の2条件を満たす非負整数列の個数をmod $ 10^{9} + 7 $で求めよ。

  • $ A_1 + A_2 + \dots + A_N = M $
  • $2A_i \leqq A_{i-1} + A_{i+1} \,\,\ (2 \leqq i \leqq N)$

制約

  • $ 1 \leqq N \leqq 10^{5}$
  • $1 \leqq M \leqq 10^5$

解法

以降狭義単調減少(増加)を単調減少(増加)と呼ぶことにします。

一旦一つ目の条件を無視して、二つ目の条件を満たす数列を考えると、(単調減少 → )一定( → 単調増加)のようになります。
例えば$ N = 4 $のとき、4 1 1 1, 1 1 1 3, 4 3 1 2は条件を満たしますが、4 1 3 2は条件を満たしません。

数列の最小値を$ C $とし、$ i $番目で最小値をとるとします。このとき$ A_i = C $です。こうすると、問題は

  • $ i $番目まで数が単調減少し、そこから$ j (i \leqq j \leqq N) $番目まで$ A_j = C$となり、$ j $番目以降から単調増加

となる数列を数え上げるということになります。ここで$ i $番目までは必ず単調減少とすることで、各$ i $について数え上げた結果に重複がないことが 保証されます。

単調減少というだけでは実は二つ目の条件を満たしません。

  • $ i $に近づくにつれて隣接する数の差が広義単調減少

を条件として追加する必要があります。

すると、例えば$ A_{i-1} = C+1 $としたとき、必ず$ i-1 $番目より左側も隣接する数の差が1以上にならなければならないため、

  • $A_{i-1} \geqq C+1$
  • $A_{i-2} \geqq C+2$

$\dots$

  • $A_1 \geqq C+(i-1)$

が成り立っています。$ C $との差分のみを足し上げると、 $$ 1 + 2 + \dots + (i-1) = \frac{1}{2} i (i-1) $$ となります。よって1つ目の条件を考えると、$ i $は、$ i \leqq \sqrt{2M} $を満たす必要があることがわかります。
以降$ m = \sqrt{2M} $とします。数列は左から順に

  • 長さ$ O(m) $の単調減少数列
  • $ O(N-m) $個の $ C $
  • 長さ$ O(m) $の単調増加数列

のようになっています。

ここで差を1増やしたら全体の和がいくつ増えるのかを考えます。

  • $ i-1 $番目の値を1増やすと、$1 + 2 + \dots + (i-1) = \frac{1}{2} i (i-1)$だけ和が増える
  • $ i-2 $番目の値を1増やすと、$1 + 2 + \dots + (i-2) = \frac{1}{2} (i-1) (i-2)$だけ和が増える

という形になります。和が$ M $になるまでどこに何回1増やしてもいいので、これは個数制限なしナップザックのような形に帰着できます。
ここで加算する候補は$ 1, 3, \dots \frac{1}{2}i(i-1) $ であり、$ O(m) $個しかないので、和が$j \,(1 \leqq j \leqq M)$となるような場合の数を求める時の計算量は$ O(M\sqrt{M})$ となります。

$$ dp[i][j] = i番目まで単調減少としたときに和がjとなるような加算方法の総数 $$ としたときの遷移を考えていきます。 ここで$ i $番目まで単調減少と決めたら、必ず$ i-1 $番目に1以上加算しなければならないと考慮すると、以下のように遷移します。

  • $0,1, \dots , i-1$番目までの$ dp[j] $の値の和を用意し($ S $とする)、$$ dp[i][j] = S + dp[i][j-\frac{1}{2}i(i+1)] $$

これで必ず1回は加算するようになります。
また、今回は1から$ m $までの全ての$ i $における和が$ j $になる場合の数が知りたいので、$ mM $個の値全てを保存します。

また、$ C $の値は0から$\frac{M}{N}$までの値をとりますが、これを愚直にやっていては間に合いません。しかし、和がちょうど$ M $ではなく、和が$ M - (Nの倍数)$ とすることで、前処理を行うことができます(具体的には$ dp[i][j+n] $ += $ dp[i][j] $とすることでできます)。

ここから$ i $を動かし、答えを求めます。
左から$ i $番目までを単調減少区間としたとき、単調増加区間の開始地点は$i, i+1, \dots, N$が考えられます。
ただし、単調増加開始地点は$ N-m $以上でないと単調増加を始められないことに注意します(和が$ M $を超えてしまうので)

左から単調増加は、右から単調減少と考えられるので、dpの値を利用することができます。
また、各$ i $について独立になるように計算されているので、右側は$ i $から$ N $までを足し合わせたものと考えることができます。
まとめると、左から$ i $番目までを単調減少とするとき、

  • 各$ j (1 \leqq j \leqq M) $について$ \displaystyle dp[i][j] * \sum_{k=0}^{k=N-i} dp[k][m-j] $を加算する

という操作をすることで答えを求められます。ただし、$ N $が大きい場合には$N-i$と$ m $のminをとります。
シグマの箇所は先ほどのdp遷移で用いた累積和を配列に格納することで、それをそのまま用いることができます。

以上で$ O(Mm) $でこの問題を解くことができます。

実装

includeやmodintの箇所は省きました。
メモリと時間はかなりギリギリです。

mint dp[450][101010], sdp[450][101010];

int main(){
    ll n,m; cin >> n >> m;
    ll sq = sqrt(m*2) + 1;
    vector<int> a(sq);
    for(int i=0; i<sq; i++) a[i] = (i+1) * (i+2) / 2;
    dp[0][0] = sdp[0][0] = 1;
    for(int i=0; i<sq; i++){
        for(int j=0; j<=m; j++){
            if(j-a[i] >= 0) dp[i+1][j] += dp[i+1][j-a[i]] + sdp[i][j-a[i]];
        }
        for(int j=0; j<=m; j++) sdp[i+1][j] += dp[i+1][j] + sdp[i][j];
    }
    for(int i=0; i<=sq; i++){
        for(int j=0; j<=m; j++) if(j+n <= m) dp[i][j+n] += dp[i][j];
    }
    mint ans = 0;
    for(int i=0; i<=min(sq,n-1); i++){
        if(n-i-1 < 0) break;
        for(int j=0; j<=m; j++){
            ans += dp[i][j] * sdp[min(n-i-1,sq)][m-j];
        }
    }
    cout << ans << "\n";
}