AOJ 2692 - ICPC Teams

問題文

http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=2692

解法

包除原理。
チームにしないといけない組は確定で組ませる。
また、チームにしてはいけない組のすべての部分集合に対して、その集合に含まれる組は確定でチームを組ませるとする。
このとき、チームの人数を3人ずつにできなくなるようなら無視する（4人以上のチームがある、2人確定したチーム数が多すぎるなど）。
すると、3人チームがいくつかと2人チームがいくつかになって、ほかは全部自由という条件になる。
自由な人から2人チームに割り当てて、あとは3人グループを作るだけ。
このときの場合の数は簡単な組み合わせ（高校数学の教科書に乗ってそうなぐらい典型なので書かないけど）で求まる。
これを最初に述べた部分集合のサイズに対して包除原理すれば求まる。

ソースコード

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

// mint (mod とるライブラリ), union_find は省略

int main() {
    int n, m; cin >> n >> m;
    n *= 3; // to number of people
    map<int, int> idx;
    vector<int> a(m), b(m), c(m), ban;
    for(int i = 0; i < m; ++i) {
        cin >> a[i] >> b[i] >> c[i];
        if(c[i] == 1) {
            ban.push_back(i);
        }
    }
    const int sz = [&] (){ // compress
        vector<int> xs;
        for(int i = 0; i < m; ++i) {
            xs.push_back(a[i]);
            xs.push_back(b[i]);
        }
        sort(begin(xs), end(xs));
        xs.erase(unique(begin(xs), end(xs)), xs.end());
        for(int i = 0; i < m; ++i) {
            a[i] = lower_bound(begin(xs), end(xs), a[i]) - begin(xs);
            b[i] = lower_bound(begin(xs), end(xs), b[i]) - begin(xs);
        }
        return xs.size();
    }();

    mint ans = 0;
    vector<mint> inv6_fact(n, mint(1));
    {
        const mint inv6 = mint(6).inverse();
        for(int i = 1; i < n; ++i) {
            inv6_fact[i] = inv6_fact[i - 1] * inv6;
        }
    }
    for(int S = 0; S < (1 << ban.size()); ++S) { // same team
        union_find uf(sz);
        for(int i = 0; i < m; ++i) {
            if(c[i] == 1) continue;
            uf.unite(a[i], b[i]);
        }
        for(int i = 0; i < (int)ban.size(); ++i) {
            if(S & (1 << i)) {
                uf.unite(a[ban[i]], b[ban[i]]);
            }
        }
        int cnt_two = 0, cnt_three = 0;
        {
            bool check = true;
            vector<bool> f(sz);
            for(int i = 0; i < sz; ++i) {
                if(f[uf.root(i)]) continue;
                f[uf.root(i)] = true;
                check &= uf.size(i) <= 3;
                cnt_three += uf.size(i) == 3;
                cnt_two += uf.size(i) == 2;
            }
            check &= (n - cnt_three * 3 - cnt_two * 2) >= cnt_two;
            if(!check) continue; // not count
        }
        int tn = n - cnt_three * 3 - cnt_two * 2;
        mint tans = comb(tn, cnt_two) * fact(cnt_two);
        tn -= cnt_two;
        tans *= fact(tn) * inv6_fact[tn / 3] * fact(tn / 3, true);
        if(__builtin_popcount(S) & 1) {
            ans -= tans;
        } else {
            ans += tans;
        }
    }

    cout << ans << endl;
}

感想

AtCoder Beginner Contest で出ても違和感ないなって感じの数え上げ典型だった。