Official

C - Filling 3x3 array Editorial by Nyaan


皆さんは今まで、パズルや数学の問題を見て「この問題はパソコンを使って計算できるんじゃないか?」と考えたことはありますか?そしていざコードを書いてみると、実際に答えが求まったり、あるいはプログラムの実行が終わらず何も得られなかったり…という経験がある方もいるのではないかと思います。
今回の問題はまさにこのようなシチュエーションを問題にしたものです。おそらく今回は、値が大きいケースで TLE してしまう人が一定数出たのではないかと思います。この問題が解けなかった人は、「適切な考察によりプログラムを高速化する」ということを学んでいただけたらと思います。

(以下の説明では \(\max(h_1,h_2,h_3,w_1,w_2,w_3) = N\) と表します。 )
まずはじめに思いつくのはナイーブな全探索です。各マスに入るのは \(1\) 以上 \(N\) 以下の整数に限られるので、実際に書くマスに \(1\) 以上 \(N\) 以下の数を代入してみて、その盤面が条件を満たすか判定する、という方法です。
これは \(9\) 重の for-loop を書けば 実装できますが少し筋が悪いので、深さ優先探索 (DFS)を用いて実装してみましょう。するとだいたい次のようになります。

int h[3], w[3], a[3][3];
long long ans = 0;
void dfs(int ij) {
  int i = ij / 3, j = ij % 3;
  if (i == 3) { 
    if( /* 盤面が条件を満たす */ ) ans++;
    return;
  }
  for (int x = 1; x <= 30; x++) {
    a[i][j] = x;
    dfs(ij + 1);
  }
}

この方針では正しい答えを得るプログラムを書くことはできますが、計算量が \(\mathrm{O}(N^9)\) になり TLE してしまうという問題があります。そこで、どうにか探索の要らない部分を省略して高速化できないか考えてみましょう。

マス目に \(a\) から \(i\) の名前を次の図のようにつけます。

image

問題の性質として、\(\bf{1}\) 行 / \(\bf{1}\) 列のうち \(\bf{2}\) 個が埋まったら残り \(\bf{1}\) 個は自動的に決まるというのがあります。
例えば、\(h_1 = 10, a = 3, b = 5\) ならば、\(c\) は次の式で求まります。

\[c = h_1 - a - b = 10 - 3 - 5 = 2\]

元々のプログラムは \(3 \times 3\) の範囲を全探索していましたが、この事実を利用すると、全探索する範囲は左上 \(\bf 2 \times 2\) だけでいい ということが示せます。

詳しく説明しましょう。探索中に \(a,b,d,e\) の値を確定させたとします。\(i\) 行目の和が \(h_i\) なのを \(i=1,2\) について利用すると \(c,f\) は次のように求まります。

\[\begin{aligned} c &= h_1 - a - b \\ f &= h_2 - d - e \\ \end{aligned} \]

さらに \(a\) から \(f\) が求まっていると、\(g,h,i\) は次のように求まります。

\[\begin{aligned} g &= w_1 - a - d \\ h &= w_2 - b - e \\ i &= w_3 - c - f \\ \end{aligned} \]

使ってない条件は \(g+h+i=h_3\) なので、これで得られた \(g,h,i\) の和が \(h_3\) と一致するかを確認すればよいです。

このように探索する範囲をうまく削減した全探索ならば、計算量は \(\mathrm{O}(N^4)\) に落ちているので、制約の範囲内では十分に高速に動作します。実装は for-loop や DFS を用いればよいです。

  • for-loop を用いた実装例 (C++)
#include <algorithm>
#include <iostream>
using namespace std;
int H[3], W[3], ans = 0;
int main() {
  for (int i = 0; i < 3; i++) cin >> H[i];
  for (int j = 0; j < 3; j++) cin >> W[j];
  for (int a = 1; a <= 30; a++) {
    for (int b = 1; b <= 30; b++) {
      for (int d = 1; d <= 30; d++) {
        for (int e = 1; e <= 30; e++) {
          int c = H[0] - a - b;
          int f = H[1] - d - e;
          int g = W[0] - a - d;
          int h = W[1] - b - e;
          int i = W[2] - c - f;
          if (min({c, f, g, h, i}) > 0 and g + h + i == H[2]) ans++;
        }
      }
    }
  }
  cout << ans << "\n";
}
  • DFS を用いた実装例 (C++)
#include <iostream>
using namespace std;
int h[3], w[3], a[3][3];
long long ans = 0;
void dfs(int ij) {
  int i = ij / 3, j = ij % 3;
  if (i == 3) {
    ans++;
    return;
  }
  if (i == 2) {
    int x = w[j] - a[0][j] - a[1][j];
    if (x <= 0) return;
    a[i][j] = x, dfs(ij + 1);
  } else if (j == 2) {
    int x = h[i] - a[i][0] - a[i][1];
    if (x <= 0) return;
    a[i][j] = x, dfs(ij + 1);
  } else {
    for (int x = 1; x <= 30; x++) a[i][j] = x, dfs(ij + 1);
  }
}
int main() {
  for (int i = 0; i < 3; i++) cin >> h[i];
  for (int j = 0; j < 3; j++) cin >> w[j];
  if (h[0] + h[1] + h[2] != w[0] + w[1] + w[2]) {
    cout << 0 << "\n";
    return 0;
  }
  dfs(0);
  cout << ans << "\n";
}

posted:
last update: