[AtCoder arc140 D] One to One (图论计数分治NTT)

145 阅读 0 评论 96 点赞

我是靠谱客的博主谦让心锁，最近开发中收集的这篇文章主要介绍[AtCoder arc140 D] One to One (图论计数分治NTT)，觉得挺不错的，现在分享给大家，希望可以做个参考。

概述

题意

初始有 $n$ 个点，给定一个长度为 $n$ 的数组 $a_i$ ，若 $a_i ne -1$ ，则有无向边 $i, a_i)$ ，若 $a_i = -1$ ，则点 $i$ 可以连向 $1 \sim n$ 任意点，求所有图的联通块个数之和

$a_i in [1, n] cup {-1}$

对 $998244353$ 取模。

分析：

首先考虑忽略 $a_i = -1$ 的所有边，那么图中会有若干个连通块，这些连通块分为三种情况：

树
基环树
环

对于环和基环树来说，因为是 $n$ 个点和 $n$ 条边，所以他们不可能有一条出边，换句话说，里边的点不可能包含 $a_i = -1$ ，而对于树来说，因为是 $n$ 个点 $n - 1$ 条边，所以有且仅有一条出边，也就是树里面只有一个 $a_i = -1$

这就代表树可以和其他连通块组成一个新的连通块，但是无论树如何连边，环和基环树的连通性都不会发生变化，也就是他始终有一个环，所以可以先计算出这部分的贡献，设图中环和基环树的数量为 $u$ ，树的数量为 $v$ ，则这部分贡献就为 $u × n v u times n ^ {v}$

接下来考虑树的所有连边情况，我们枚举 $k$ 条边组成一个环，设第 $i$ 棵树的大小为 $f_i$ ，每棵树则有生成函数 $1 + f_ix$ ，记 $F (x)$ 为选若干个树构成一个环的方案数，可以用分治 $NTT$ 快速求出。
$prod_{i = 1} ^ {v} (1 + f_ix)$
每个点构成一个 $k$ 元环是有顺序的，第一个点可以有 $k - 1$ 种选择，第二个点有 $k - 2$ 种选择，所以总共构成一个 $k$ 元环的方案数为 $(k - 1)!$ ，还要考虑剩下没有被选出来的点，那么可以随便连，都不影响这个环，方案数就为 $n ^ {v - k}$ ，那么答案就是
$sum_{k = 1} ^ {v} (k - 1)! times [x ^ k]F(x) times n ^ {v - k}$
时间复杂度 $O(nlog ^ 2n)$

代码：

#include <bits/stdc++.h>
#define int long long
#define poly vector<int>
#define len(x) ((int)x.size())
using namespace std;
const int N = 3e3 + 5, G = 3, Ginv = 332748118, mod = 998244353;
struct DSU {
    vector<int> p, Size;
    DSU(int n) : p(n), Size(n, 1) {
        iota(p.begin(), p.end(), 0);
    }
    int find(int x) {
        return p[x] == x ? p[x] : p[x] = find(p[x]);
    }
    bool same(int u, int v) {
        return find(u) == find(v);
    }
    void merge(int u, int v) {
        u = find(u), v = find(v);
        if (u != v) {
            Size[v] += Size[u];
            p[u] = v;
        }
    }
};
int lim;
vector<int> rev(N);
int qmi(int a, int b) {
    int res = 1;
    while (b) {
        if (b & 1) res = res * a % mod;
        a = a * a % mod;
        b >>= 1;
    }
    return res;
}
void polyinit(int n) {
    for (lim = 1; lim < n; lim <<= 1);
    for (int i = 0; i < lim; i ++) rev[i] = (rev[i >> 1] >> 1) | (i & 1 ? lim >> 1 : 0);
}
void NTT(poly &f, int op) {
    for (int i = 0; i < lim; i ++) {
        if (i < rev[i]) swap(f[i], f[rev[i]]);
    }
    for (int mid = 1; mid < lim; mid <<= 1) {
        int Gn = qmi(op == 1 ? G : Ginv, (mod - 1) / (mid << 1));
        for (int i = 0; i < lim; i += mid * 2) {
            for (int j = 0, G0 = 1; j < mid; j ++, G0 = G0 * Gn % mod) {
                int x = f[i + j], y = G0 * f[i + j + mid] % mod;
                f[i + j] = (x + y) % mod, f[i + j + mid] = (x - y + mod) % mod;
            }
        }
    }
    if (op == -1) {
        int inv = qmi(lim, mod - 2);
        for (int i = 0; i < lim; i ++) f[i] = f[i] * inv % mod;
    }
}
poly operator * (poly f, poly g) {
    int n = len(f) + len(g) - 1;
    polyinit(n), f.resize(lim), g.resize(lim);
    NTT(f, 1), NTT(g, 1);
    for (int i = 0; i < lim; i ++) f[i] = f[i] * g[i] % mod;
    NTT(f, -1), f.resize(n);
    return f;
}
signed main() {
    cin.tie(0) -> sync_with_stdio(0);
    int n;
    cin >> n;
    DSU p(n + 1);
    vector<int> st(n + 1), fact(n + 1);
    fact[0] = 1;
    for (int i = 1; i <= n; i ++) {
        fact[i] = fact[i - 1] * i % mod;
    }
    for (int i = 1; i <= n; i ++) {
        int x;
        cin >> x;
        if (x != -1) {
            int flag = 0;
            if (p.same(i, x)) {
                flag = 1;
            }
            p.merge(i, x);
            if (flag) {
                st[p.find(i)] = 1;
            }
        }
    }
    int v = 0, u = 0;
    vector<poly> f(n + 1);
    for (int i = 1; i <= n; i ++) {
        if (p.find(i) == i) {
            if (!st[i]) {
                v ++;
                f[v].resize(2);
                f[v][0] = 1, f[v][1] = p.Size[i];
            } else {
                u ++;
            }
        }
    }
    int res = u * qmi(n, v) % mod;
    if (v) {
        function<poly(int, int)> dc = [&](int l, int r) {
            if (l == r) {
                return f[l];
            }
            int mid = l + r >> 1;
            return dc(l, mid) * dc(mid + 1, r);
        };
        poly ans = dc(1, v);
        for (int k = 1; k <= v; k ++) {
            res = (res + fact[k - 1] * ans[k] % mod * qmi(n, v - k) % mod) % mod;
        }
    }
    cout << res << endl;
}