2023-07-17

AtCoder ABC 310 D - Peaceful Teams (水色, 400 点)

最小添字規則によってダブルカウントを防ぐ対象を一意に定める操作列を数え上げる AtCoder AtCoder400点水色diff ABC-D 全探索グラフグルーピンググルーピングの数え上げ問題数え上げ問題全探索:再帰関数 DP bitDP O(3^N) 順列の数え上げ見積り大事制約条件:共演NG O(3^N)なbitDP 典型要素を詰め合わせた教育的問題【問題集】ビット全探索が困難な再帰全探索問題【問題集】DFS・BFSのステップアップ

再帰関数を使った全探索が書けるかどうかが試される！

問題へのリンク

問題概要

$N$ 人のスポーツ選手を $T$ 個のグループに分けたい。

ただし、仲の悪いスポーツ選手の組が $M$ 組ある。任意の $i = 0, 1, \dots, M-1$ に対して、選手 $A_{i}$ と選手 $B_{i}$ を同じグループに属させてはならない。

この制約を守りつつ、 $T$ 個のグループに分ける方法は何通りあるか求めよ。

制約

$1 \le T \le N \le 10$
$0 \le M \le \frac{N(N-1)}{2}$

再帰関数を用いた全探索の復習

こういう全探索問題に直面するとき、人はビット全探索を再帰関数でも書けるかどうかが試されるのであろう。

たとえば、 $N$ 個の品物の中からいくつか選ぶ方法のうち、条件を満たすものが何通りあるかを調べたいとする。このとき、ビット全探索なら次のように書ける。

int res = 0;
for (int bit = 0; bit < (1<<N); ++bit) {
    // bit のビットが立っている桁を集める
    vector<int> a;
    for (int i = 0; i < N; ++i) {
        if (bit & (1<<i)) a.push_back(i);
    }

    // 選ぶ方法 a が条件に合うかを調べる
    if (a が条件に合う) ++res;
}

しかし、世の中には、綺麗に $2^{N}$ 通りを調べれば済むような全探索ではない全探索もたくさんある！

それに備えて、ビット全探索と同等の処理を、ぜひ再帰関数でも書けるようにしておきたい (次の記事を参照)。

drken1215.hatenablog.com

上のようなビット全探索は、たとえば、次のように書ける。

// 品物 0, 1, ..., i-1 のうち、すでに選んだ品物の集合が a と表されて
// 今後 i 番目以降の品物を選ぶかどうかを検討する場面からスタートした場合に、
// 条件を満たすものが何通りあるか
int dfs(vector<int> &a, int i) {
    // 終端条件
    if (i == N) {
       if (a が条件に合う) return 1;
       else return 0;
    }

    // 答えを表す変数
    int res = 0;

    // a に i を追加しない場合
    res += dfs(a, i + 1);
    
    // a に i を追加する場合
    a.push_back(i);
    res += dfs(a, i + 1);
    a.pop_back();  // 元に戻す (バックトラック)

    return res;
}

この再帰関数に対して、空の vector<int> 型変数 a を用いて、dfs(a, 0) を呼び出すことで、答えが求められる。

なお、a.pop_back(); などというように、「元に戻す」処理は忘れがちなので注意。このように、元に戻す処理を絡めながら再帰関数を回す探索を、バックトラックなどと呼んだりする。

今回の問題

上記のような再帰関数が書けるなら、今回の問題もさほど難しくないはず。

今度は次のような再帰関数を用意しよう！

// 選手 0, 1, ..., i-1 のうち、すでに作られたグループが groups と表されるとき、
// 今後、選手 i 以降をどのグループに入れるかを検討する場面からスタートした場合に、
// 条件を満たすものが何通りあるか
int dfs(vector<vector<int>> &groups, int i) {
    // 終端条件
    if (i == N) {
       if (groups が条件に合う) return 1;
       else return 0;
    }

    // 答えを表す変数
    int res = 0;
        
    // 新しいグループを作って、そこに選手 i を入れる場合を試す
    groups.push_back(vector<int>({i}));
    res += dfs(groups, i + 1);
    groups.pop_back();    // 元に戻す

    // 既存グループに選手 i を入れていく方法を試していく
    for (auto &g : groups) {
        if (グループ g に選手 i を入れられる) {
            g.push_back(i);
            res += dfs(groups, i + 1);
            g.pop_back();   // 元に戻す
        }
    }
    return res;
}

あとは細部を詰めることで、AC できる！

最後に、計算時間を簡単に見積もろう。一般に、 $i = 0, 1, \dots, N-1$ 人目の選手をグループに入れる方法は最悪でも $i+1$ 通りである。よって、調べるべき場合の数は、グループ数が $T$ 個以下という制約を無視したとしても、 $N!$ 通りにおさまることがわかる。

また、再帰関数の各ノードにおける探索処理は $O(N)$ の計算量となる。よって、探索全体の計算量は $O(N N!)$ と評価できる。

コード

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

// 入力
int N, T, M;
bool ng[11][11];

int dfs(vector<vector<int>> &groups, int i) {
    // 終端条件
    if (i == N) {
        // サイズが T であるならば、条件を満たす
        if (groups.size() == T) return 1;
        else return 0;
    }

    int res = 0;
    
    // 新しいグループを作って、そこに選手 i を入れる場合を試す
    groups.push_back(vector<int>({i}));
    res += dfs(groups, i+1);
    groups.pop_back();  // 元に戻す

    // 既存グループに選手 i を入れていく方法を試していく
    for (auto &g : groups) {
        // groups 内の各グループ g について、選手 i を入れるとダメな場合はスキップ
        bool ok = true;
        for (auto v : g) if (ng[v][i]) ok = false;
        if (!ok) continue;
        
        // グループ g に選手 i を入れる場合を試す
        g.push_back(i);
        res += dfs(groups, i+1);
        g.pop_back();  // 元に戻す
    }
    return res;
}

int main() {
    cin >> N >> T >> M;
    for (int i = 0; i < M; ++i) {
        int A, B;
        cin >> A >> B;
        --A, --B;
        ng[A][B] = ng[B][A] = true;
    }

    vector<vector<int>> groups;
    cout << dfs(groups, 0) << endl;
}