けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ARC 062 E - AtCoDeerくんと立方体づくり / Building Cubes with AtCoDeer (橙色, 900 点)

大変だった ^^;

問題へのリンク

問題概要

全部で  C 色の色がある。
下図のような  N 枚の正方形があって、それぞれ  1, 2, \dots, N のラベルがついていて、各正方形の四隅にはいずれかの色が塗られている。

 N 枚の中から 6 枚を選んで立方体を作る。ただしどの頂点についても、そこに集まっている 3 個の色が等しくなければならない。

そのようなものが何通りあるかを数えよ。ただし各面の「1」のような数値ラベルは回転すると異なる模様になるので、回転して色の配置が等しくなるとしても異なるものとして数える。しかし最終的に出来上がる立方体を回転すればまったく同じ「ラベル」「色付け」になるものを同じものとして数える。

制約

  •  1 \le N \le 400
  •  1 \le C \le 1000

考えたこと

昨日会社の同僚・後輩たちと一緒に解いた!

 N \le 400 くらいの制約の問題に苦手意識がある...
さて、回転をまともに考えるのは嫌なので

  • 上面を先に固定してしまう ( N 通り)
  • 下面は回転も含めて固定 ( 4(N-1) 通り)

という風にすることにした。そうすると側面がどうあるべきかが決まるので原理的には集計できるはず。愚直な実装でも  O(N^{2} \log{N}) でできそう ( \log は std::map とか使うため)。注意点として、立方体の一番上の面にどれを持ってくるかで 6 の重複度があるので最後に 6 で割ることにする。

ここで

  • map<vector, long long> count // 色合いが vector(4 次元) で表される長方形が何通りあるか

という連想配列を用意しておくことにする。上面を固定したときに、それをすでに使ったことを表現するために

  • count から上面の色合いを c (4 次元ベクトル) として、c の回転も含めた 4 パターンをそれぞれ count から引いておく
  • 探索の終端で足す

という風にすれば OK。

以上の数え上げでよいことを、サンプル 3 (全部の色が 0) の意味を解読することで確認する。

  • 上面を 6 通り
  • 下面を残り 5 個から選んで回転: 5 × 4 通り
  • 各側面の選び方がそれぞれ 4 × 4、3 × 4、2 × 4、1 × 4 通り

これを全部掛け算して 6 で割ると確かに一致する。

#include <iostream>
#include <vector>
#include <map>
#include <algorithm>
using namespace std;

// 各色の長方形が何個ずつあるか
map<vector<int>, long long> ma;
void use(vector<int> v) {
    for (int it = 0; it < 4; ++it) {
        vector<int> col(4);
        for (int j = 0; j < 4; ++j) {
            col[j] = v[(it+j)%4];
        }
        ma[col]--;
    }
}
void add(vector<int> v) {
    for (int it = 0; it < 4; ++it) {
        vector<int> col(4);
        for (int j = 0; j < 4; ++j) {
            col[j] = v[(it+j)%4];
        }
        ma[col]++;
    }
}

int main() {
    int N; cin >> N;
    vector<vector<int> > C(N, vector<int>(4));
    for (int i = 0; i < N; ++i) for (int j = 0; j < 4; ++j) cin >> C[i][j];

    // 追加
    ma.clear();
    for (int i = 0; i < N; ++i) add(C[i]);

    // 集計
    long long res = 0;
    for (int i = 0; i < N; ++i) {
        use(C[i]);
        for (int j = 0; j < N; ++j) {
            if (j == i) continue;
            use(C[j]);
            
            for (int dir = 0; dir < 4; ++dir) {
                auto ue = C[i];
                vector<int> shita(4);
                for (int k = 0; k < 4; ++k) shita[k] = C[j][(dir + k) % 4];
                reverse(shita.begin(), shita.end());

                long long tmp = 1;
                vector<bool> isadd(4, 0);
                for (int it = 0; it < 4; ++it) {
                    vector<int> col = {ue[(it+1)%4], ue[it%4], shita[it%4], shita[(it+1)%4]};
                    if (!ma.count(col)) tmp = 0;
                    else tmp *= ma[col], use(col), isadd[it] = true;
                }
                for (int it = 0; it < 4; ++it) {
                    vector<int> col = {ue[(it+1)%4], ue[it%4], shita[it%4], shita[(it+1)%4]};
                    if (isadd[it]) add(col);
                }
                
                res += tmp;
            }
            add(C[j]);
        }
        add(C[i]);
    }
    cout << res/6 << endl;
}