Official

G - Row Column Sums 2 Editorial by leaf1415


\(\sum_{i = 1}^N R_i \neq \sum_{i = 1}^N C_i\) のときは明らかに答えは \(0\) です。 以下、\(\sum_{i = 1}^N R_i = \sum_{i = 1}^N C_i\) を仮定します。

\(1\) 行目、\(2\) 行目、\(\ldots\)\(N\) 行目の順で一行ずつ埋めていき、問題文中の条件を満たす行列 \(A\) を作ることを考えます。 \(i\) 行目を埋める際は、\(A_{i, 1}, A_{i, 2}, \ldots, A_{i_N}\) を、

  • \(\sum_{j = 1}^N A_{i, j} = R_i\) かつ、
  • どの列 \(j\) においても \(C'_{i, j} := C_j - \sum_{k = 1}^i A_{k, j} \geq 0\)

であるように決定します。 \(C'_{i, j}\)\(i\) 行目まで埋めたときの \(j\) 列目の 残りの和 と言えます。 \(C'_{i, j}\) の値は \(0, 1, 2\) のいずれかであることに注意してください。 \(N\) 行目すべてを埋める方法であって、すべての列 \(j\)\(C'_{N, j} = 0\) となるものの個数が本問題の答えです。

動的計画法(DP)で本問題を解きます。DP テーブルとして、

\(\mathrm{dp}[i][x] =\)\(A\)\(i\) 行目までを埋める方法であって、\(C'_{i, j} = 2\) であるような列 \(j\) の個数が \(x\) 個であるものの個数)

とおきます。

\(A\)\(i\) 行目まで埋めたとき、\(C'_{i, j}= 2\) である列 \(j\) の個数が \(x\) 個であることがわかれば、このときの \(C'_{i, j} = 1\) である列 \(j\) の個数 \(y\) も、入力で与えられた定数と \(x\) から求まります。 実際、

\[\sum_{j = 1}^N C'_{i, j} = y + 2x\]

が成り立つので、

\[ \begin{aligned} y &= \sum_{j = 1}^N C'_{i, j} -2x\\ &= \sum_{j = 1}^N (C_j - \sum_{k = 1}^i A_{k, j}) -2x\\ &= \sum_{j = 1}^N C_j - \sum_{j = 1}^N \sum_{k = 1}^i A_{k, j} - 2x\\ &= \sum_{j = 1}^N C_j - \sum_{k = 1}^i \sum_{j = 1}^N A_{k, j} - 2x\\ &= \sum_{j = 1}^N C_j - \sum_{k = 1}^i R_k - 2x   \tag{1} \end{aligned} \]

です。ここで、最後の等号には \(\sum_{j = 1}^N A_{i, j} = R_i\) を満たすように各 \(i\) 行目を埋めていることを用いました。

それでは、具体的な DP の方法について述べます。 まず、DP の初期状態を考えます。 これは、\(C'_{0, j} = C_j = 2\) を満たす列 \(j\) の個数を \(x_0\) とおくと、

\[ \mathrm{dp}[0][x] = \begin{cases} 1 &(x = x_0)\\ 0 &(x \neq x_0) \end{cases} \]

です。

次に、DP の遷移を考えます。 \(\mathrm{dp}[i-1][\ast]\) が既知のとき、\(\mathrm{dp}[i][x]\) を求めることを考えます。 \(A\)\(i\) 行目まで埋めた結果、\(C'_{i, j} = 2\) である列 \(j\) の個数が \(x\) であったとします。まず、\(C'_{i, j} = 1\) である列 \(j\) の個数 \(y\) を式 (1) で求めます。 そして、\(R_i\) の値によって以下の \(3\) 通りに場合分けします。

(i) \(R_i = 0\) のとき

\(i\) 行目の埋め方として下記の \(1\) 通りのみが考えられます。

  • \(C'_{i-1, j} = 1, 2\) である列 \(j\) がそれぞれ \(x\), \(y\) 個ある状態から、\(i\) 行目をすべて \(0\) で埋めた。

よって、

\[\mathrm{dp}[i][x] = \mathrm{dp}[i-1][x]\]

です。

(ii) \(R_i = 1\)のとき

\(i\) 行目の埋め方として下記の \(2\) 通りが考えられます。

  • \(C'_{i-1, j} = 1, 2\) である列 \(j\) がそれぞれ \(x+1\), \(y-1\) 個ある状態から、\(C'_{i-1} = 2\) である一列 \(j_1\) について \(A_{i, j_1} = 1\) とし、残りを \(0\) で埋めた。
  • \(C'_{i-1, j} = 1, 2\) である列 \(j\) がそれぞれ \(x\), \(y+1\) 個ある状態から、\(C'_{i-1} = 1\) をである一列 \(j_1\) について \(A_{i, j_1} = 1\) とし、残りを \(0\) で埋めた。

よって、

\[\mathrm{dp}[i][x] = \mathrm{dp}[i-1][x+1]\cdot (x+1) + \mathrm{dp}[i-1][x] \cdot (y+1)\]

です。

(iii) \(R_i = 2\) のとき

\(i\) 行目の埋め方として下記の \(4\) 通りが考えられます。

  • \(C'_{i-1, j} = 1, 2\) である列 \(j\) がそれぞれ \(x+1\), \(y\) 個ある状態から、\(C'_{i-1} = 2\) である一列 \(j_1\) について \(A_{i, j_1} = 2\) とし、残りを \(0\) で埋めた。
  • \(C'_{i-1, j} = 1, 2\) である列 \(j\) がそれぞれ \(x+2\), \(y-2\) 個ある状態から、\(C'_{i-1} = 2\) である二列 \(j_1, j_2\) について \(A_{i, j_1} = A_{i, j_2} = 1\) とし、残りを \(0\) で埋めた。
  • \(C'_{i-1, j} = 1, 2\) である列 \(j\) がそれぞれ \(x\), \(y+2\) 個ある状態から、\(C'_{i-1} = 1\) である二列 \(j_1, j_2\) について \(A_{i, j_1} = A_{i, j_2} = 1\) とし、残りを \(0\) で埋めた。
  • \(C'_{i-1, j} = 1, 2\) である列 \(j\) がそれぞれ \(x+1\), \(y\) 個ある状態から、\(C'_{i-1} = 1\) である一列 \(j_1\)\(C'_{i-1} = 2\) を満たす \(1\)\(j_2\) について \(A_{i, j_1} = A_{i, j_2} = 1\) とし、残りを \(0\) で埋めた。

よって、

\[\mathrm{dp}[i][x] = \mathrm{dp}[i-1][x+1]\cdot (x+1) + \mathrm{dp}[i-1][x+2]\cdot \binom{x+2}{2} + \mathrm{dp}[i-1][x]\cdot \binom{y+2}{2} + \mathrm{dp}[i-1][x+1]\cdot (x+1)y\]

です。

以上の初期化と遷移によって、DPテーブルを埋めることができます。 \(\sum_{i = 1}^N R_i = \sum_{i = 1}^N C_i\) という冒頭の仮定より、\(C'_{N, j} = 2\) である列 \(j\) の個数 \(x\)\(0\) のとき、\(C'_{N, j} = 1\) である列 \(j\) の個数 \(y\) も式 (1) より \(0\) となるので、本問題の答えは \(\mathrm{dp}[N][0]\) として求まります。

以下に C++ 言語による正解例を記載します。

#include <iostream>
#include <atcoder/modint>
using namespace std;
using namespace atcoder;
using mint = modint998244353;

int n;
int r[5005], c[5005];
mint dp[5005][5005];

int main(void){
  cin >> n;
  for(int i = 1; i <= n; i++) cin >> r[i];
  for(int i = 1; i <= n; i++) cin >> c[i];
  
  int x0 = 0, sum = 0;
  for(int i = 1; i <= n; i++){
    sum += c[i];
    if(c[i] == 2) x0++;
  }
  
  dp[0][x0] = 1;
  for(int i = 1; i <= n; i++){
    sum -= r[i];
    for(int x = 0; x <= n; x++){
      int y = sum - 2*x;
      if(y < 0) break;
      if(r[i] == 0) dp[i][x] = dp[i-1][x];
      if(r[i] == 1) dp[i][x] = dp[i-1][x+1]*(x+1) + dp[i-1][x]*(y+1);
      if(r[i] == 2) dp[i][x] = dp[i-1][x+1]*(x+1) + dp[i-1][x+2]*((x+2)*(x+1)/2) + dp[i-1][x]*((y+2)*(y+1)/2) + dp[i-1][x+1]*(x+1)*y;
    }
  }
  
  if(sum) cout << 0 << endl;
  else cout << dp[n][0].val() << endl;
  
  return 0;
}

posted:
last update: