AtCoder Regular Contest (ARC) 124 E - Pass to Next の解法です． コンテスト後にも解けませんでした． 公式解説 を読んだだけではわからず， laycrsさんのツイートとコード を読んで，なんとか理解しました．

問題へのリンク

https://atcoder.jp/contests/arc124/tasks/arc124_e

問題概要

1, 2, …, N の番号のついた人が円環状に並ぶ． 各人は a_i 個のボールを持っている． 0個以上 a_i 個のボールを右隣の人に一斉に渡す (1回限り)． その結果 人 i が持つボールの数を x_i とする． 有り得る $(x_i \mid i \in [1, N])$ 全体の集合を S とする． $\sum \{ \prod \{ x_i \mid i \in [1, N] \} \mid x \in S \}$ を mod 998244353 で求めよ．

解法

(わかってしまえば，公式解説と書いていることは同じです)

異なる渡し方が同じ $x$ を生成するための条件は，渡す個数の差が 全員同じであることであるから，「渡す個数が 0 の人が一人以上いる」 という条件をつけると，生成される $S$ には変化がなく，かつ，渡し方と $S$ が1対1に対応する．

$x$ を一つ固定すると，$\prod \{ x_i \mid i \in [1, N] \}$ は， 各人の操作後のボールから1つを選ぶ選び方の数と一致する．

求める答は，以下の塗り方の数と一致する:

  • 各 $i \in [1, N] $ に対し，$a_i + 1$ 枚のカードが1列に横に並んで グループをなしており， 全体では円環上に並んでいる (グループ N の次にグループ 1 が来る)．
  • 以下の条件を満たすように，いくつかのカードを黒く， いくつかのカードを赤く塗る．
    • 各グループのカードのうち，ちょうど1枚を黒く塗る．
    • 黒く塗るカードが右端のカードであるグループが1つ以上有る．
    • どの2枚の黒いカードの間にも，赤いカードがちょうど1枚有る．

各グループの，黒いカードより右側のカードが，右隣の人に渡したボールに 対応し，黒いカードより右側のカードは，手元に残したボールに対応する． 赤いカードは，「選んだ」ボールに対応する．

このような塗り方を，以下の DP で数える．

  • dp[i][a][b][c] := 第 $i$ グループまでを以下の条件を満たすように塗る方法の数
    • $a \in \{0, 1\}$． $a = 0$ では，最後のグループで，黒いカードより右側に赤いカードが無いことを 仮定する． $a = 1$ では，有ることを仮定する．
    • $b \in \{0, 1\}$． $b = 0$ は，第 $i$ グループで，黒いカードよりも右に 赤いカードが無いことを示す． $b = 1$ は，有ることを示す．
    • $c \in \{0, 1\}$． $c = 0$ は，第 1 グループから第 $i$ グループまででは， 右端のカードを黒く塗ったものが無いことを示す． $c = 1$ は，有ることを示す．

遷移は下のコードを参照．

求める答は，dp[N][0][0][1] + dp[N][1][1][1] である．

ACコード

Fp binom(ll n, ll r) {   // 二項係数．3までしか出てこない
  if (r < 0 || n < r) return 0;
  if (r == 0) return 1;
  if (r == 1) return n;
  if (r == 2) return Fp(n) * Fp(n - 1) / Fp(2);
  if (r == 3) return Fp(n) * Fp(n - 1) * Fp(n - 2) / Fp(6);
  assert(0);
}

int main(/* int argc, char *argv[] */) {
  ios_base::sync_with_stdio(false);
  cin.tie(nullptr);
  cout << setprecision(20);

  ll N; cin >> N;
  vector<ll> A(N);
  for (ll i = 0; i < N; i++) cin >> A[i];

  auto tbl_init = vector(2, vector(2, vector(2, Fp(0))));
  auto tbl = tbl_init;
  tbl[0][0][0] = tbl[1][1][0] = 1;
  for (ll i = 0; i < N; i++) {
    auto prev = move(tbl);
    tbl = tbl_init;
    for (ll a = 0; a < 2; a++) {             // 遷移前後でaは変わらず
      for (ll b = 0; b < 2; b++) {           // 遷移前
        for (ll b0 = 0; b0 < 2; b0++) {      // 遷移後
          for (ll c = 0; c < 2; c++) {       // 遷移前
            for (ll c0 = 0; c0 < 2; c0++) {  // 遷移後
              // n 枚から r 枚を選んで，下の (*) で更新する．
              ll n = A[i] + 1;
              // r は最大で 黒1, 赤2 の 3枚だが，
              // グループ i-1 で選択されていたら1減じる．
              // グループ i で選択しないのであれば1減じる．
              ll r = 3 - b - (1 - b0);
              if (c && !c0) continue;    // 矛盾
              else if (!c && c0) {
                // このグループで右端を黒く塗る
                if (b0) continue;        // 右端より右には置けない
                // 右端を黒に決めたので，n, r とも1減じる．
                n--;
                r--;
              }else if (!c && !c0 && !b0) n--; // 右端は選べない
              tbl[a][b0][c0] += binom(n, r) * prev[a][b][c];     // (*)
            }
          }
        }
      }
    }
  }
  cout << tbl[0][0][1] + tbl[1][1][1] << endl;

  return 0;
}

Fp は，mod 998244353 を計算するための構造体です． https://atcoder.jp/contests/arc124/submissions/25361161