問題文

出典 ARC120 F

方針

まず，N本のワインからK本連続しないように取る総数は$\binom{N-K+1}{K}$ですね．これは選ぶワインをoとし，選ばれなかったワインをxとすると，先頭K-1本はoxをセットで考えると分かりますね．これを$f(N, K)$と書きます．

すると，i番目を取る総数はfの畳み込みの形で表されます．FFTなどを考えますが，恐らく2変数の畳み込みが必要でかなり厳しい気持ちになります．二項係数のsumはキレイな形の場合二項係数になる場合があります．が，今回sumの添字が中途半端なのでうまくいきません．多分．なのでこの方針から離れることにします．

初等的な組み合わせ論の方針を考えます．必ず取るものがi番目で，それが真ん中のあたりにある時，

……xox……

みたいな形になりますね．……の部分に着目すると，左側の……の一番右と右側の……の一番左は同時にoでも良いことを除くと元の問題とそんなに変わっていない気がします．これはxox部分をxに変えたものの総数と等しいです．

……x……

これはどうやって数えればよいかというと，真ん中のxの位置に何が来てもいいものの総数から真ん中の位置にoがあるものを引けばいいですね．

……(o or x)…… - ……o……

この左の項はN-2個からK-1個のoを選ぶ問題で，右の項はN-2個からi-1番目を必ず取るようにしてK-1個選ぶ総数になります．

N個あってi番目を選ぶという条件のもとでK個取る総数を$g(N, i, K)$と書くと

$$ g(N, i, K) = f(N-2, K-1) - g(N-2, i-1, K-1) $$

みたいな式になります．fはiに関わらず同じなのでなんかうまくいきそうな気配がしますね．

最初にiが真ん中のあたりにあるって条件を書きましたが，もし$i=0, 1$のとき

$$g(N, 0, K) = f(N-2, K-1)$$

$$g(N, 1, K) = f(N-3, K-1)$$

になるので注意が必要です．あと，iは簡単のため端からの距離として定義しておきます．これは左右対称なのでOKです．iは$\min(i, N-1-i)$とすれば良いです．

ソースの中ではfを交互に正負を逆転させ累積和を取る関数をhとし，$i=1$に到達した時の正負を含めて計算する関数をlとしています．lはオーバーしているかの判定をしていないですが，オーバーしている時は二項係数が0になるのでうまくいきます．

ソースコード

#include<cstdio>
using namespace std;

typedef long long ll;
const ll NCK_MAX = 510000;
const ll mod = 998244353;

ll *fac, *finv, *inv;

void nCk_init(ll mod) {
  fac = new ll[NCK_MAX];
  finv = new ll[NCK_MAX];
  inv = new ll[NCK_MAX];
  fac[0] = fac[1] = 1;
  finv[0] = finv[1] = 1;
  inv[1] = 1;
  for (ll i = 2; i < NCK_MAX; i++) {
    fac[i] = fac[i-1] * i % mod;
    inv[i] = mod - inv[mod%i] * (mod / i) % mod;
    finv[i] = finv[i-1] * inv[i] % mod;
  }
}

ll nCk(ll n, ll k, ll mod) {
  if (fac == NULL) nCk_init(mod);
  if (n < k) return 0;
  if (n < 0 || k < 0) return 0;
  return fac[n] * (finv[k] * finv[n-k] % mod) % mod;
}

ll n, k, d, a[300000], dp[300001], ans;

ll f(ll n, ll i) {
  return nCk(n-i+1, i, mod);
}

ll h(ll i) {
  if (i <= 1) return 0;
  if (dp[i] != 0) return dp[i];
  ll tmp = f(n-2*(i-1), k-(i-1));
  if (i % 2 == 1) tmp = mod - tmp;
  return dp[i] = (tmp + h(i-1)) % mod;
}

ll l(ll x) {
  ll tmp = f(n-2*x-1, k-x);
  if (x % 2 == 0) tmp = mod - tmp;
  return tmp;
}

ll g(ll x) {
  return h(x) + l(x);
}


int main() {
  scanf("%lld", &n);
  scanf("%lld%lld", &k, &d);
  for (ll i = 0; i < n; i++) scanf("%lld", &a[i]);
  ll sum = 0;
  for (ll i = 0; i < n; i++) {
    ll x = 2 * i < n ? i : n - 1 - i;
    if (x == 0) {
      ans += f(n-2, k-1) * a[i] % mod;
    } else {
      ans += g(x) * a[i] % mod;
    }
  }
  printf("%lld\n", ans % mod);
}

感想

無限に頭がバグります．解法ガチャでいい方針を引かないと地獄だと思いました．

コンテスト中に通したかったなぁ…