29. 牛客-一人行者

最新推荐文章于 2024-10-09 22:28:58 发布

theophania

最新推荐文章于 2024-10-09 22:28:58 发布

阅读量145

点赞数

分类专栏：一题文章标签：算法动态规划

本文链接：https://blog.csdn.net/theophania/article/details/126674827

版权

一题专栏收录该内容

41 篇文章 0 订阅

订阅专栏

博主分享了在解决牛客网竞赛题目102D中遇到的问题，即树形动态规划中关于联通子集计数时，如何避免除以零的错误。通过调整策略，记录前缀积以防止dp[1][u]为mod-1导致的除法问题，从而找到正确答案。

摘要由CSDN通过智能技术生成

本来不想为了这题写一篇博客的，但因为昨天被一组数据卡了一个小时，还是有必要来记录一下。

牛客练习赛 102D：一人行者

题意是给一棵树，求断掉每一条边后得到的两棵树各自的联通子集数量，对 $998244353$ 取模。

容易想到树形 dp，令 $d p [u] [0/1]$ 表示 $u$ 的子树中是否包含 $u$ 的子集数量，显然有状态转移方程：

$\left\{ \begin{array}{l} dp[u][0]=\sum (dp[v][0]+dp[v][1])\\ dp[u][1]=\prod (1+dp[v][1]) \end{array} \right.$

其中 $v$ 是 $u$ 的子节点。

然后考虑断边的操作，假设断的边是 $(u, v)$ ，其中 $v$ 是儿子，那么 $v$ 所在的子树的答案 $V = d p [v] [0] + d p [v] [1]$ ，而 $u$ 所在的树的答案是总的联通子集个数减去 $V$ 再减去以 $v$ 作为中间节点的联通子集个数。那么还需要求出每个节点“向上的子树”中联通子集的个数，其实就是换根 dp，重新进行一次 dfs 后自上而下计算即可。

于是我交了一份这样的代码：

#include <bits/stdc++.h>
using namespace std;
using std::cin;
using std::cout;
using std::string;
#define endl '\n'
#define TRACE(x) std::cerr << #x << " = " << x << endl
#define fi first
#define se second
#define pb push_back

using uint = unsigned int;
using ll = long long;
using ull = unsigned long long;
using pii = std::pair<int, int>;

constexpr int maxn = 5e5 + 5;
constexpr int inf = 0x3f3f3f3f;
constexpr ll mod = 998244353;
template <uint MOD = mod> struct mint {
    uint x;
    mint() : x(0) {}
    mint(ll _x) {
        _x %= MOD;
        if (_x < 0)
            _x += MOD;
        x = _x;
    }
    explicit operator int() const { return x; } 
    mint &operator+=(const mint &a) {
        x += a.x;
        if (x >= MOD)
            x -= MOD;
        return *this;
    }
    mint &operator-=(const mint &a) {
        x += MOD - a.x;
        if (x >= MOD)
            x -= MOD;
        return *this;
    }
    mint &operator*=(const mint &a) {
        x = (ull)x * a.x % MOD;
        return *this;
    }
    mint pow(ll pw) const {
        mint res = 1;
        mint cur = *this;
        while (pw) {
            if (pw & 1)
                res *= cur;
            cur *= cur;
            pw >>= 1;
        }
        return res;
    }
    mint inv() const {
        assert(x != 0);
        uint t = x;
        uint res = 1;
        while (t != 1) {
            uint z = MOD / t;
            res = (ull)res * (MOD - z) % MOD;
            t = MOD - t * z;
        }
        return res;
    }
    mint &operator/=(const mint &a) { return *this *= a.inv(); }
    mint operator+(const mint &a) const { return mint(*this) += a; }
    mint operator-(const mint &a) const { return mint(*this) -= a; }
    mint operator*(const mint &a) const { return mint(*this) *= a; }
    mint operator/(const mint &a) const { return mint(*this) /= a; }
    bool operator==(const mint &a) const { return x == a.x; }
    bool operator!=(const mint &a) const { return x != a.x; }
    bool operator<(const mint &a) const { return x < a.x; }
    friend std::istream &operator>>(std::istream &is, mint &a) {
        ll v;
        is >> v;
        a = mint(v);
        return is;
    }
    friend std::ostream &operator<<(std::ostream &os, const mint &a) {
        return os << a.x;
    }
};

int fa[maxn];
pii q[maxn];
vector<int> g[maxn];
using Z = mint<998244353>;
Z dp[2][maxn];
void dfs(int u, int f) {
    fa[u] = f;
    dp[1][u] = 1;
    for (auto& v : g[u]) {
        if (v == f) continue;
        dfs(v, u);
        dp[0][u] += dp[0][v] + dp[1][v];
        dp[1][u] *= (dp[1][v] + 1);
    }
}
Z up[maxn];
void dfs1(int u, int f) {
    if (f) {
        up[u] = (up[f] + 1) * dp[1][f] / (dp[1][u] + 1);
    }
    for (auto& v : g[u]) {
        if (v == f) continue;
        dfs1(v, u);
    }
}
void solve() {
    int n;
    cin >> n;
    for (int i = 1, u, v; i < n; ++i) {
        cin >> u >> v;
        g[u].pb(v), g[v].pb(u);
        q[i] = {u, v};
    }
    dfs(1, 0);
    dfs1(1, 0);
    for (int i = 1; i < n; ++i) {
        auto [u, v] = q[i];
        bool f = 0;
        if (fa[u] == v) {
            f = 1;
            swap(u, v);
        }
        Z V = dp[0][v] + dp[1][v];
        Z U = dp[0][1] + dp[1][1] - V - up[v] * dp[1][v];
        if (f) swap(U, V);
        cout << U << " " << V << endl;
    }
}

int main() {
    std::ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int T = 1;
    // cin >> T;
    for (int cas = 1; cas <= T; ++cas) {
        solve();
    }
    return 0;
}