ABC294-Ex 包除原理

問題

Ex - K-Coloring
$N$頂点$M$辺の無向グラフ$G=(V,E)$が与えられる。 各頂点を$1\sim K$の色で塗る方法であって、隣り合った頂点が異なる色で塗られているものの個数を求めよ。

解説

解説放送では包除原理があっさり説明されていたのでその部分を書く。\(\{ 1,2,\dots,N \}\)を$[N]$と書く。
彩色 $ c: V \to [K] $全体の集合を$U$とする。 $i$番目の辺$(u_i,v_i)$について、集合$A_i$を \(A_i = \{ c \in U \mid c(u_i) = c(v_i) \}\) と定義すると、求める答えは \(\left| \bigcap _{i=1}^M \overline{A_i} \right|\) になる。$\left| \bigcap _{i\in S} \overline{A_i} \right|$は難しいが、$ \left| \bigcap _{i\in S} A_i \right| $は$i( \in S)$番目の辺を残した部分グラフの連結成分数を$N_S$とすると$K^{N_S}$になるため簡単に求まる。 そのため包除原理

\[\left| \bigcup_{i \in U} B_i \right| = \sum_{\emptyset \ne S \subseteq U} (-1)^{|S|+1} \left| \bigcap_{i \in S} B_i \right|\]

を使って変形する。

\[\left| \bigcap _{i \in [M]}\overline{A_i} \right| = \overline{\left| \bigcup _{i \in [M]} A_i\right|} = \left|U\right|- \left| \bigcup _{i \in [M]} {A_i} \right| = \left|U\right| - \sum_{\emptyset \ne S \subseteq [M]} (-1)^{|S|+1} \left| \bigcap_{i \in S} {A_i} \right|\]

したがって

\[\left| \bigcap _{i \in [M]}\overline{A_i} \right| = \sum_{S \subseteq [M]} (-1)^{|S|} \left| \bigcap_{i \in S} {A_i} \right| = \sum_{S \subseteq [M]} (-1)^{|S|} K^{N_S}\]

となる。下の形の式のほうがよく使われるかもしれない。

コード(TLE)

#include "/home/y_midori/cp/library/template.hpp"
#include <atcoder/dsu>
#include <atcoder/modint>
using mint = atcoder::modint998244353;
void solve() {
    LL(n, m, k);
    vector<P> edge(m);
    rep(i, m) {
        LL(a, b);
        a--, b--;
        edge[i] = P(a, b);
    }
    mint ans = 0;
    for(unsigned long long bit = 0; bit < 1 << m; bit++) {
        atcoder::dsu g(n);
        rep(i, m) {
            if(bit >> i & 1) {
                auto [a, b] = edge[i];
                g.merge(a, b);
            }
        }
        ans += (popcount(bit) & 1 ? -1 : 1) * mint(k).pow(g.groups().size());
        debug(ans.val());
    }
    print(ans.val());
}

int main() {
    ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    cout << std::setprecision(16);
    int t = 1;
    rep(_, t) {
        solve();
    }
}

(本題終わり)
これでは$S \subseteq [M]$の列挙に$2^M \le 2^{30}$回かかり間に合わないので工夫する必要がある。 頂点$i$をグラフから取り除いたとき、次数が$0$のとき取り除いたグラフの答えの$K$倍、$1$のときは隣り合う頂点と別の色に塗れば良いから$K-1$倍になることがすぐ分かる。そうでない場合簡単には頂点を取り除けないから、$i$と繋がっている$1$つの辺を取り除くことを考える。繋がっている先の頂点を$v$とすると、答えは辺を取り除いたときの答えから$i,v$が同じ色としたときの答えを引いたものとなる。この場合後者は$i$と$v$を同一視し、$i$と繋がっていた辺を$v$に繋ぎ変え、多重辺を取り除くことで$i$を取り除いて計算することができる。以上のことを再帰的に計算する

コード(AC)

#include "/home/y_midori/cp/library/template.hpp"
#include <atcoder/modint>
using mint = atcoder::modint998244353;
void solve() {
    LL(n, m, k);
    vector<P> edge(m);
    rep(i, m) {
        LL(a, b);
        a--, b--;
        if(a > b)
            swap(a, b);
        edge[i] = P(a, b);
    }
    auto dfs = [&k](auto dfs, ll n, vector<P> edge) -> mint {
        if(n == 0)
            return 1;
        vl deg(n);
        for(auto [a, b] : edge) {
            deg[a]++, deg[b]++;
        }
        vl ord(n);
        iota(all(ord), 0);
        sort(all(ord), [&deg](ll i, ll j) { return deg[i] > deg[j]; });
        vl rev(n);
        rep(i, n) rev[ord[i]] = i;
        for(auto &[a, b] : edge) {
            a = rev[a], b = rev[b];
            if(a > b)
                swap(a, b);
        }
        sort(rbegin(deg), rend(deg));
        {
            mint res = 1;
            while(n and deg.back() == 0) {
                res *= k;
                n--;
                deg.pop_back();
            }
            if(n == 0 or res != 1) {
                return res * dfs(dfs, n, edge);
            }
        }
        ll idx = -1, v = -1;
        rep(i, edge.size()) {
            debug(edge[i].first, edge[i].second);
            if(edge[i].second == n - 1) {
                idx = i;
                v = edge[i].first;
                break;
            }
        }
        edge.erase(begin(edge) + idx);
        if(deg.back() == 1) {
            return dfs(dfs, n - 1, edge) * (k - 1);
        }
        mint res = dfs(dfs, n, edge);
        for(auto &[a, b] : edge) {
            if(b == n - 1) {
                b = v;
                if(a > b)
                    swap(a, b);
            }
        }
        sort(all(edge));
        edge.erase(unique(all(edge)), edge.end());
        return res - dfs(dfs, n - 1, edge);
    };
    print(dfs(dfs, n, edge).val());
}

int main() {
    ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    cout << std::setprecision(16);
    int t = 1;
    rep(_, t) {
        solve();
    }
}