🐷

【ABC256】「C - Filling 3x3 array」のメモ

2022/07/18に公開

概要

東京海上日動プログラミングコンテスト2022(AtCoder Beginner Contest 256)でC問題が解けなかったので、解説を見て得られた知見などをメモしておきます。
https://atcoder.jp/contests/abc256/tasks/abc256_c

コンテスト中に書いたコード

マス目の数が9個ありますので、それぞれについて全探索してチェックすれば一応解けます。
計算量的にアウトなのは明らかでしたが、計算量を削減する方法がわからなかったので一応そのまま実装してみました。狂気の9重ループです。

2次元配列で扱ってもよかったのですが、ちょっとわかりづらいですが1次元配列でも扱えるのでそのようにしています。

#include <bits/stdc++.h>
using namespace std;

int main() {
    int h1, h2, h3, w1, w2, w3;
    cin >> h1 >> h2 >> h3 >> w1 >> w2 >> w3;

    vector<int> table(9);

    int mx = max(h1, max(h2, (max(h3, max(w1, max(w2, w3))))));

    int ans = 0;
    for(int a = 1; a < mx; ++a) {
        table[0] = a;
        for(int b = 1; b < mx; ++b) {
            table[1] = b;
            for(int c = 1; c < mx; ++c) {
                table[2] = c;
                for(int d = 1; d < mx; ++d) {
                    table[3] = d;
                    for(int e = 1; e < mx; ++e) {
                        table[4] = e;
                        for(int f = 1; f < mx; ++f) {
                            table[5] = f;
                            for(int g = 1; g < mx; ++g) {
                                table[6] = g;
                                for(int h = 1; h < mx; ++h) {
                                    table[7] = h;
                                    for(int i = 1; i < mx; ++i) {
                                        table[8] = i;
                                        if(table[0] + table[1] + table[2] == h1 &&
                                            table[3] + table[4] + table[5] == h2 &&
                                            table[6] + table[7] + table[8] == h3 &&
                                            table[0] + table[3] + table[6] == w1 &&
                                            table[1] + table[4] + table[7] == w2 &&
                                            table[2] + table[5] + table[8] == w3) {
                                                ans++;
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    cout << ans << endl;
}

解説を読んで

自分が全然知らない解法が解説に載っているんだろうなと思っていたのですがそんなことはなく、ベースの考え方自体は同じで、「2つ決まったらもう一つは決まるのでその分は列挙しなくていい」という性質を利用して探索範囲を狭めるというアプローチになっていました。
それなら頑張って考えれば解けていたかもしれないので悔しいですね。
解説を見て修正したコードは以下です。

#include <bits/stdc++.h>
using namespace std;

int main() {
    int h1, h2, h3, w1, w2, w3;
    cin >> h1 >> h2 >> h3 >> w1 >> w2 >> w3;

    vector<int> table(9);

    int mx = max(h1, max(h2, (max(h3, max(w1, max(w2, w3))))));

    int ans = 0;
    for(int a = 1; a < mx; ++a) {
        table[0] = a;
        for(int b = 1; b < mx; ++b) {
            table[1] = b;
            table[2] = h1 - table[0] - table[1];

            if(table[2] < 1) {
                continue;
            }

            for(int c = 1; c < mx; ++c) {
                table[3] = c;
                for(int d = 1; d < mx; ++d) {
                    table[4] = d;
                    table[5] = h2 - table[3] - table[4];
                    table[6] = w1 - table[0] - table[3];
                    table[7] = w2 - table[1] - table[4];
                    table[8] = h3 - table[6] - table[7];


                    if(table[5] < 1 || table[6] < 1 || table[7] < 1 || table[8] < 1) {
                        continue;
                    }

                    if(table[0] + table[1] + table[2] == h1 &&
                        table[3] + table[4] + table[5] == h2 &&
                        table[6] + table[7] + table[8] == h3 &&
                        table[0] + table[3] + table[6] == w1 &&
                        table[1] + table[4] + table[7] == w2 &&
                        table[2] + table[5] + table[8] == w3) {
                            ans++;
                    }
                }
            }
        }
    }

    cout << ans << endl;
}

ループカウンタの値で計算するとどれがどのマス目の値かわかりづらいので、一旦配列に入れて計算しています。
table[0]table[1] が決まれば table[2] が、さらに table[3]table[4] が決まれば残りは全て計算で求められるのでそのようにしています。
ただし、マス目には正の整数しか書かないので、0以下になったときはスキップします。

模範解答ほどスマートじゃないですが一応ACです。

教訓

  • 範囲を絞った全探索という考え方がある
  • 特定の考え方を知らなきゃ解けないとは限らず、自分が考えたことの一歩先に答えがある場合もある

追記

DFSで書き直してみました。(言語が変わっているのはさておき)

var h1 = 0
var h2 = 0
var h3 = 0
var w1 = 0
var w2 = 0
var w3 = 0

var count = 0

fun main() {
    val input = readLine()!!.split(" ").map { it.toInt() }
    h1 = input[0]
    h2 = input[1]
    h3 = input[2]
    w1 = input[3]
    w2 = input[4]
    w3 = input[5]

    dfs(emptyList())

    println(count)
}

private fun dfs(list: List<Int>) {
    if(list.size == 4) {
        // println(list)
        val m1 = list[0]
        val m2 = list[1]
        val m3 = h1 - m1 - m2
        val m4 = list[2]
        val m5 = list[3]
        val m6 = h2 - m4 - m5
        val m7 = w1 - m1 - m4
        val m8 = w2 - m2 - m5
        val m9 = h3 - m7 - m8

        if(listOf(m3, m6, m7, m8, m9).any { it < 1 }) {
            return
        }

        if(m1 + m2 + m3 == h1 &&
            m4 + m5 + m6 == h2 &&
                m7 + m8 + m9 == h3 &&
                m1 + m4 + m7 == w1 &&
                m2 + m5 + m8 == w2 &&
                m3 + m6 + m9 == w3) {
            count++
        }
        return
    }

    for(i in 1..30) {
        dfs(list + listOf(i))
    }
}

Discussion