CF 1499F Diameter Cuts

题目描述

https://codeforces.com/contest/1499/problem/F

Solution

经典的树形 $dp$，我们令 $f[u][i]$ 表示 $u$ 子树内都合法且从 $u$ 向下延伸的最长路径长度为 $i$

由于 $i$ 一定小于等于 $sz[u]$，所以转移可以做到 $O(n^2)$

时间复杂度 $O(n^2)$

#include <iostream>
#include <cstring>
#include <vector>
#include <queue>
#include <algorithm>
#define maxn 5010
#define ll long long
using namespace std;

const int p = 998244353;

int n, k;

struct Edge {
    int to, next;
} e[maxn * 2]; int c1, head[maxn];
inline void add_edge(int u, int v) {
    e[c1].to = v; e[c1].next = head[u]; head[u] = c1++;
}

int sz[maxn]; ll f[maxn][maxn], g[maxn];
void dfs(int u, int fa) {
    sz[u] = 1; f[u][0] = 1; 
    for (int i = head[u]; ~i; i = e[i].next) {
        int v = e[i].to; if (v == fa) continue;
        dfs(v, u); 
        for (int j = 0; j <= min(::k, sz[u]); ++j)
            for (int k = 0; k <= min(sz[v], ::k - j); ++k)
                g[max(j, k)] = (g[max(j, k)] + f[u][j] * f[v][k]) % p;
        sz[u] += sz[v];
        for (int j = 0; j <= min(::k, sz[u]); ++j) f[u][j] = g[j], g[j] = 0;
    } ll sum = 0; 
    for (int i = 0; i <= min(k, sz[u]); ++i) sum = (sum + f[u][i]) % p;
    if (u != 1) {
        for (int i = min(k, sz[u]); i; --i) f[u][i] = f[u][i - 1];
        f[u][0] = sum;
    }
} 

int main() { fill(head, head + maxn, -1); 
    ios::sync_with_stdio(false);
    cin.tie(nullptr); cout.tie(nullptr);

    cin >> n >> k;
    for (int i = 1; i < n; ++i) {
        int x, y; cin >> x >> y;
        add_edge(x, y); add_edge(y, x);
    } dfs(1, 0);
    ll ans = 0; for (int i = 0; i <= k; ++i) ans = (ans + f[1][i]) % p;
    
    cout << ans << "\n";
    return 0;
}