2021 ICPC Southeastern Europe Regional Contest C Werewolves

题目描述

https://codeforces.com/gym/103438/problem/C

简要题意：给定一棵 $n$ 个点的树，每个点有一个颜色 $c_i$，求有多少连通块满足存在某种颜色的出现次数大于其它所有颜色的出现次数的和

$n\le 3000$

Solution

我们首先考虑枚举颜色，把是这个颜色点赋值成 $1$，不是变成 $-1$，那么我们跑一个树形背包，容易发现大于 $0$ 的都是合法的

这样时间复杂度看起来是 $O(n^3)$ 的，但注意到如果我令当前枚举的颜色为 $c$，该颜色的总数为 $k_c$，那么如果背包的大小下雨 $-k_c$ 是没有意义的，所以我们可以把背包总容量限制为 $[-k_c,k_c]$，那么这样做一次的复杂度就是 $O(nk_c)$，所以总的时间复杂度为 $O(n\sum k_c=n^2)$

#include <iostream>
#include <cstring>
#include <vector>
#define maxn 3010
#define ll long long
using namespace std;

const int offset = 3000;
const int p = 998244353;

int n, a[maxn], c[maxn];

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 * 2], g[maxn * 2], ans;
void dfs(int u, int fa, int L) {
    sz[u] = 1; for (int i = -L; i <= L; ++i) f[u][i + offset] = 0;
    f[u][a[u] + offset] = 1;
    for (int i = head[u]; ~i; i = e[i].next) {
        int v = e[i].to; if (v == fa) continue;
        dfs(v, u, L);
        for (int S = -L; S <= L; ++S) g[S + offset] = 0;
        for (int S = -min(sz[u], L); S <= min(sz[u], L); ++S)
            for (int T = -min(sz[v], L); T <= min(sz[v], L); ++T)
                g[max(S + T, -L) + offset] = (g[max(S + T, -L) + offset] + f[u][S + offset] * f[v][T + offset]) % p;
        sz[u] += sz[v];
        for (int S = -L; S <= L; ++S) f[u][S + offset] = g[S + offset];
    } ++f[u][0 + offset];
    for (int i = 1; i <= L; ++i) ans = (ans + f[u][i + offset]) % p;
}

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

    cin >> n;
    for (int i = 1; i <= n; ++i) cin >> c[i];
    for (int i = 1; i < n; ++i) {
        int x, y; cin >> x >> y;
        add_edge(x, y); add_edge(y, x); 
    }
    for (int col = 1; col <= n; ++col) {
        int cnt = 0; 
        for (int i = 1; i <= n; ++i) a[i] = c[i] == col ? 1 : -1, cnt += a[i] == 1;
        dfs(1, 0, cnt); 
    } cout << ans << "\n";
    return 0;
}