牛客2022多校DAY10-K You are given a tree

「牛客2022多校DAY10-K」 You are given a tree...

简要题意

给一棵带点权和边权的树，找到至多 $k$ 个点权不同的点，使得它们之间路径覆盖的边权和最大。

$n\le 5000,k\le 5$。

Solution

考虑颜色数量不大的时候怎么暴力。显然可以直接状压 DP，压一下当前选择的颜色集合为 $S$，在子树 $i$ 内的最大覆盖边权为 $f(i,S)$，转移的时候 $\mathcal{O}(3^n)$ 枚举子集，可以做到 $\mathcal{O}(n{3}^n)$ 的复杂度。

注意到 $k$ 并不大，所以选择的颜色也只有很少，如果能让颜色数量也变得跟 $k$ 一样少就行了，因此想到随机化。考虑将每个颜色随机映射到 $[1,k]$ 的一个数上，然后使用上面的做法。这种做法显然会出错，因此计算一下出错的概率：

原来的 $k$ 个颜色在随机之后互不相同的概率为 ${\displaystyle \frac{k!}{k^k}}$，当 $k=5$ 时大约是 $0.038$，那么错误率为 $0.962$，因此重复计算 $200$ 次仍然出错的概率为 ${0.962}^{200}\approx 0.0004$，正确率大概是 $0.9996$，除非脸太黑，否则都可以过。

Code

#include<bits/stdc++.h>
#define int long long

using namespace std;

namespace Hanx16qwq {
constexpr int _N = 5e3 + 5;
int n, k, a[_N], b[_N], col[_N];
struct EDGE{int nxt, to, l;} edge[_N << 1];
int tot, head[_N];
void AddEdge(int x, int y, int l) {
    edge[++tot] = {head[x], y, l};
    head[x] = tot;
}
int f[_N][64], temp[64], maxn, ans = INT_MIN;
void DfsForAns(int x, int fa) {
    f[x][0] = f[x][1 << col[x]] = 0;
    for (int i = head[x]; i; i = edge[i].nxt) {
        int twd = edge[i].to;
        if (twd == fa) continue;
        DfsForAns(twd, x);
        memcpy(temp, f[x], sizeof f[x]);
        for (int s = 1; s < maxn; s++)
            for (int ss = s; ss; ss = (ss - 1) & s)
                if (ss != (maxn - 1))
                    f[x][s] = max(f[x][s], temp[s ^ ss] + f[twd][ss] + edge[i].l);
    }
    ans = max(ans, f[x][maxn - 1]);
}
mt19937 rnd(114514);
void main() {
    cin >> n >> k; maxn = 1 << k;
    for (int i = 1; i <= n; i++) cin >> a[i];
    for (int i = 1, u, v, l; i < n; i++) {
        cin >> u >> v >> l;
        AddEdge(u, v, l);
        AddEdge(v, u, l);
    }
    int Times = 200;
    while (Times--) {
        for (int i = 1; i <= n; i++) b[i] = rnd() % k;
        for (int i = 1; i <= n; i++) col[i] = b[a[i]];
        memset(f, 0xcf, sizeof f);
        DfsForAns(1, 0);
    }
    cout << ans << '\n';
}

}
signed main() {
    cin.tie(0)->sync_with_stdio(0);
    Hanx16qwq::main();
    return 0;
}