AGC049 D : Convex Sequence
問題リンク
問題概要
整数$ 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"; }