AOJ 2958 Tree (HUPC 2019 day2-G)

二乗の木 DP を応用したもの。このセットの運営チームだった。

• 問題へのリンク
• editorial
• 北大アーカイブ

問題概要

頂点からなる木が与えられる。

この木の 個の空でない部分グラフの組であって、各部分グラフがいずれも連結で、どの二つの相異なる部分グラフも頂点を共有しないものの個数を 998244353 で割ったあまりを求めよ。

制約

• 
• 

考えたこと

これをする！

snuke.hatenablog.com

いつも通り根付き木にして

• dp[ v ][ k ].first := v を根とする根付き木において、頂点 v を使わない場合について、k 個の部分グラフをとる方法の個数
• dp[ v ][ k ].second := v を根とする根付き木において、頂点 v を使う場合について、k 個の部分グラフをとる方法の個数

とする。これらを普通に木 DP で計算すると の計算量となるが、sunke さんの記事の方法を用いることで となる。

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

// modint
template<int MOD> struct Fp {
    long long val;
    constexpr Fp(long long v = 0) noexcept : val(v % MOD) {
        if (val < 0) val += MOD;
    }
    constexpr int getmod() const { return MOD; }
    constexpr Fp operator - () const noexcept {
        return val ? MOD - val : 0;
    }
    constexpr Fp operator + (const Fp& r) const noexcept { return Fp(*this) += r; }
    constexpr Fp operator - (const Fp& r) const noexcept { return Fp(*this) -= r; }
    constexpr Fp operator * (const Fp& r) const noexcept { return Fp(*this) *= r; }
    constexpr Fp operator / (const Fp& r) const noexcept { return Fp(*this) /= r; }
    constexpr Fp& operator += (const Fp& r) noexcept {
        val += r.val;
        if (val >= MOD) val -= MOD;
        return *this;
    }
    constexpr Fp& operator -= (const Fp& r) noexcept {
        val -= r.val;
        if (val < 0) val += MOD;
        return *this;
    }
    constexpr Fp& operator *= (const Fp& r) noexcept {
        val = val * r.val % MOD;
        return *this;
    }
    constexpr Fp& operator /= (const Fp& r) noexcept {
        long long a = r.val, b = MOD, u = 1, v = 0;
        while (b) {
            long long t = a / b;
            a -= t * b, swap(a, b);
            u -= t * v, swap(u, v);
        }
        val = val * u % MOD;
        if (val < 0) val += MOD;
        return *this;
    }
    constexpr bool operator == (const Fp& r) const noexcept {
        return this->val == r.val;
    }
    constexpr bool operator != (const Fp& r) const noexcept {
        return this->val != r.val;
    }
    constexpr bool operator < (const Fp& r) const noexcept {
        return this->val < r.val;
    }
    friend constexpr istream& operator >> (istream &is, Fp<MOD>& x) noexcept {
        is >> x.val;
        x.val %= MOD;
        if (x.val < 0) x.val += MOD;
        return is;
    }
    friend constexpr ostream& operator << (ostream &os, const Fp<MOD>& x) noexcept {
        return os << x.val;
    }
    friend constexpr Fp<MOD> modpow(const Fp<MOD> &a, long long n) noexcept {
        if (n == 0) return 1;
        auto t = modpow(a, n / 2);
        t = t * t;
        if (n & 1) t = t * a;
        return t;
    }
};

const int MOD = 998244353;
using mint = Fp<MOD>;
using pint = pair<mint,mint>;
using Graph = vector<vector<int>>;

int N, K;
vector<vector<pint>> dp; // (v nonuse, v use)
void rec(const Graph &G, int v, int p = -1) {
    for (auto c : G[v]) {
        if (c == p) continue;
        rec(G, c, v);
    }
    vector<pint> dp2(2), ndp2;
    dp2[0] = {1, 0}, dp2[1] = {0, 1};
    for (auto c : G[v]) {
        if (c == p) continue;
        int S = min(K+1, (int)(dp2.size()+dp[c].size()));
        ndp2.assign(S, pint(0, 0));
        for (int i = 0; i < dp2.size(); ++i) {
            for (int j = 0; j < dp[c].size() && i+j-1 <= K; ++j) {
                if (i+j <= K) {
                    ndp2[i+j].first += dp2[i].first * (dp[c][j].first + dp[c][j].second);
                    ndp2[i+j].second += dp2[i].second * (dp[c][j].first + dp[c][j].second);
                }
                if (i+j-1>=0) {
                    ndp2[i+j-1].second += dp2[i].second * dp[c][j].second;
                }
            }
        }
        swap(dp2, ndp2);
    }
    dp[v] = dp2;
}

int main() {
    cin >> N >> K;
    Graph G(N);
    for (int i = 0; i < N-1; ++i) {
        int x, y; cin >> x >> y; --x, --y;
        G[x].push_back(y), G[y].push_back(x);
    }
    dp.assign(N, vector<pint>());
    rec(G, 0);
    cout << dp[0][K].first + dp[0][K].second << endl;
}