2022 年牛客多校第十场 K 题题解

K You are given a tree...

题意：给定一个 $nn$ 个节点的树，树上每个点有颜色 $a_{i}a\_i$，边有边权。问从中选择 $kk$ 个颜色不同的点，他们构成的生成子图的边权值和最大值。$n\le 1\times 10^{3}n\; \backslash leq\; 1\backslash times\; 10^3$，$k\in [2,5]k\; \backslash in\; [2,5]$，$a_i\le 1\times 10^{9}a\_i\; \backslash leq\; 1\backslash times\; 10^9$。

解法：若颜色种类非常少（小于等于 $1010$），可以直接使用状压 dp：$f_{u,S}f\_\{u,S\}$ 表示在子树 $uu$ 内选择了颜色集合为 $SS$ 的最大边权值和，枚举子集，使用子树合并，则有：

这样的时间复杂度为 $\mathcal{O}(n{3}^k)\backslash mathcal\; O(n3^k)$。

但是本题的颜色范围过大，而 $kk$ 又非常小，因而大部分颜色其实根本用不到选不了，因而可以当成同一种颜色看待。可以考虑使用随机化，把所有的颜色归入到 $kk$ 类中，然后使用上述树形 dp 求解得到答案。若颜色种类有 $mm$ 种（$m\ge km\; \backslash geq\; k$），则正确概率可以估计为 $p={\displaystyle \frac{1}{\left(\genfrac{}{}{0px}{}{m}{k}\right)}}p=\backslash dfrac\{1\}\{\{m\; \backslash choose\; k\}\}$，随机 $200200$ 次正确概率已经相当大了。

因而总的时间复杂度为 $\mathcal{O}(Tn{3}^k)\backslash mathcal\; O(Tn3^k)$。

#include <bits/stdc++.h>
using namespace std;
const int N = 5000;
struct line
{
    int from, to;
    long long v;
    int next;
};
struct line que[2 * N + 5];
int cnt, headers[N + 5];
void add(int from, int to, long long v)
{
    cnt++;
    que[cnt].from = from;
    que[cnt].to = to;
    que[cnt].v = v;
    que[cnt].next = headers[from];
    headers[from] = cnt;
}
long long f[N + 5][1 << 5];
int a[N + 5], col[N + 5], n, k;
long long ans;
void dfs(int place, int father)
{
    f[place][0] = f[place][1 << col[a[place]]] = 0;
    for (int i = headers[place]; i; i = que[i].next)
        if (que[i].to != father)
        {
            dfs(que[i].to, place);
            for (int S = 1; S < 1 << k;S++)
                f[que[i].to][S] += que[i].v;
            for (int S = (1 << k) - 1; S >= 0; S--)
            {
                for (int T = S; T; T = (T - 1) & S)
                {
                    f[place][S] = max(f[place][S], f[place][S ^ T] + f[que[i].to][T]);
                    if (S ^ T)
                        ans = max(ans, f[place][S ^ T] + f[que[i].to][T]);
                }
            }
        }
}
int main()
{
    unsigned seed = std::chrono::system_clock::now().time_since_epoch().count();
    mt19937 rand_num(seed);
    scanf("%d%d", &n, &k);
    for (int i = 1; i <= n;i++)
        scanf("%d", &a[i]);
    uniform_int_distribution<int> dist(0, k - 1);
    for (int i = 1, u, v, w; i < n; i++)
    {
        scanf("%d%d%d", &u, &v, &w);
        add(u, v, w);
        add(v, u, w);
    }
    long long res = 0;
    for (int times = 1; times <= 200; times++)
    {
        memset(f, 0xcf, sizeof(f));
        ans = 0;
        for (int i = 1; i <= n; i++)
            col[i] = dist(rand_num);
        dfs(1, 1);
        res = max(res, ans);
    }
    printf("%lld", res);
    return 0;
}