ECR151-E Boxes and Balls - zawatinはてなブログ

問題

https://codeforces.com/contest/1845/problem/E (diff: 2500)

$N$ 個の箱が横一列に並んでおり、各箱は空 $a_{i} = 0$ またはボールが一個入っている $a_{i} = 1$ 。あなたは以下の操作を丁度 $K$ 回行う。

操作: ボールが入っている箱と空箱を一つ選ぶ。これらは隣接している必要がある。選んだ箱に入っているボールを空箱に移す。

操作後の箱の状態についてあり得る通り数 $\mod 1000000007$ を求めよ。

制約

$2\ \le\ N\ \le\ 1500$
$1\ \le\ K\ \le\ 1500$

解法

以下のようなdpを考える。

$\text{dp}_{i, j, k}$ 前から $i$ 個目の箱を考慮した、現在 $i$ 個目の箱に「左からこの箱を通過しようとしているボールの数 - 右から既にこの箱を通過したボールの数」の値が $j$ 、その状態に達する為に必要な最小の操作回数が $k$ 回であるような通り数 *1

解は $K$ が奇数なら $\text{dp}_{n, 0, 1} + \text{dp}_{n, 0, 3} + \cdots + \text{dp}_{n, 0, K}$ 、偶数なら $\text{dp}_{n, 0, 0} + \text{dp}_{n, 0, 2} + \cdots + \text{dp}_{n, 0, K}$

dpの遷移は以下の様になる

$i$ 個目の箱にボールを置く

$j + a_{i}$ が正のときはそのまま流れてきたボールを一個そこで捨てることになる。 $j + a_{i} - 1$ 個のボールを $i + 1$ 番目の箱へ流すので操作回数が $j + a_{i} - 1$ 回増える。

$j + a_{i}$ が $0$ 以下のときは右から追加で一個のボールを流してくる必要がある。 $i + 1$ 番目の箱からの移動の寄与 $|j + a_{i} - 1|$ 分操作回数が増える。

結局 $\text{dp}_{i + 1, j + a_{i} - 1, k + |j + a_{i} - 1|}$ に $\text{dp}_{i, j, k}$ を配ることになる。

$i$ 個目の箱にボールを置かない

置く場合と同様に考えると、 $\text{dp}_{i + 1, j + a_{i}, k + |j + a_{i}|}$ に $\text{dp}_{i, j, k}$ を配ることになる。

普通にやると $O(N^{2}K)$ である。しかし、以下のようなポイントがある。

$j$ は $i \to i + 1$ で高々 $1$ しか増減しない
$k$ が $K$ を超えるような遷移は無視してよい
$1 + 2 + \cdots + \sqrt{K}$ が大体 $K$

このことから $j$ の添字を $O(\sqrt{K})$ 程度に減らすことができて、 $O(NK^{\frac{3}{2}})$ でできる。

実装

#include <bits/stdc++.h>

#include "Src/Template/IOSetting.hpp"
using namespace zawa;

#include "atcoder/modint"
using mint = atcoder::modint1000000007;

int main() {
    SetFastIO();

    int N, K; std::cin >> N >> K;
    constexpr int S{101}, O{50};
    std::vector dp(S, std::vector<mint>(K + 1));
    dp[O][0] = mint::raw(1);
    for (int i{} ; i < N ; i++) {
        int a; std::cin >> a;
        std::vector next(S, std::vector<mint>(K + 1));
        for (int j{} ; j < S ; j++) {
            for (int k{} ; k <= K ; k++) {
                if (dp[j][k] == mint{}) continue;
                {
                    int nj{j - 1 + a};
                    int nk{k + std::abs(nj - O)};
                    if (0 <= nj and nj < S and nk <= K) {
                        next[nj][nk] += dp[j][k];
                    }
                }
                {
                    int nj{j + a};
                    int nk{k + std::abs(nj - O)};
                    if (0 <= nj and nj < S and nk <= K) {
                        next[nj][nk] += dp[j][k];
                    }
                }
            }
        }
        dp = std::move(next);
    }
    mint ans{};
    for (int k{K % 2} ; k <= K ; k += 2) ans += dp[O][k];
    std::cout << ans.val() << '\n';
}

感想

バチャで解けず、解説AC
自分の考察はかなりダメダメだった。そもそも最小操作回数が $K$ 回になるようなものだけ数えようとしていた。阿呆。
既に置いた $1$ の数に注目する定石の手法をベースに考察していたが、今回はうまくいかないようだ。
ナップザックでよく現れる種類数がルートで抑えられるやつに似ているが、 $j \gt \sqrt{K}$ を考慮しなくて良いことが違う。
$3$ 乗が通るらしい。面白い問題なのにもったいない。

*1:実際には両方の方向からボールが通過するような状況はありえないことに注意