原题题面：https://ac.nowcoder.com/acm/contest/33195/K
题目大意
给定包含 n ( 1 ≤ n ≤ 5000 ) n(1\le n\le 5000) n(1≤n≤5000)个点的树 T = ( V , E ) T=(V,E) T=(V,E)，节点 i i i具有权值 a i ( 1 ≤ a i ≤ n ) a_i(1\le a_i\le n) ai(1≤ai≤n)，每条边具有一定边权值 w i ( − 1 0 9 ≤ w i ≤ 1 0 9 ) w_i(-10^9\le w_i\le10^9) wi(−109≤wi≤109)。
求子集 S ⊆ V S\subseteq V S⊆V，满足最多有 k ( 2 ≤ k ≤ 5 ) k(2\le k\le5) k(2≤k≤5)个点具有不同的权值，且最大化“生成边”的权值和。边 e e e是生成边，当且仅当存在 u , v ∈ S u,v\in S u,v∈S，边 e e e在 u , v u,v u,v的简单路径上。
输出最大的“生成边”的权值和。
解题思路
考虑若点权值种数较少，可以采用树状数组和状压 d p dp dp，设 d p x , V dp_{x,V} dpx,V表示 x x x节点的子树中，已选点权种类二进制为 V V V的最大边权值，枚举儿子 x x x，枚举 x x x的方案 S S S，则其他儿子的方案为 V − S V-S V−S，进行 d p dp dp转移即可，设点全集为 A A A。
由此可得转移式：
d p x , V = m a x ( d p x , V , d p x , V − S + d p v , S + w s ) ( V ⊆ A , S ⊆ V , s ∈ s o n x ) dp_{x,V}=max(dp_{x,V},dp_{x,V-S}+dp_{v,S}+w_s)(V\subseteq A,S\subseteq V,s\in son_x) dpx,V=max(dpx,V,dpx,V−S+dpv,S+ws)(V⊆A,S⊆V,s∈sonx)
但本题权值最多有 n n n种，需枚举 x , v , S , V x,v,S,V x,v,S,V，总复杂度为 O ( N 3 n ) O(N3^n) O(N3n)，显然不行，可以用随机化优化，将每个点的权值处理为 [ 1 , k ] [1,k] [1,k]，从而降低复杂度，处理200次，可将正确率升至 99.96 % 99.96\% 99.96%，时间复杂度 O ( T N 3 k ) O(TN3^k) O(TN3k)，最多约大于 2 × 1 0 8 2\times10^8 2×108，勉强可过。
代码实现
#include<bits/stdc++.h>
#define ll long long
using namespace std;
template <class T> inline void read(T&x){
char c,f=' ';
while(!isdigit(c=getchar()))f=c;
x=c^48;
while(isdigit(c=getchar()))x=(x<<3)+(x<<1)+(c^48);
if(f=='-')x=-x;
}
struct node{
int v;
ll w;
node(int v,int w):v(v),w(w){}
};
const int N=5e3+5,T=300;
const ll inf=0x3f3f3f3f3f3f3f;
vector<node>ve[N];
int n,k,a[N],b[N],c[N];
ll ans,f[N][1<<5],t;
void dfs(int u,int fa){
for (int S=0;S<1<<k;++S)f[u][S]=-inf;
f[u][0]=0;
f[u][1<<b[u]]=0;
for(int i=0;i<ve[u].size();i++){
int v=ve[u][i].v;
if(v==fa)continue;
dfs(v,u);
for(int V=0;V<1<<k;V++)f[v][V]+=ve[u][i].w;
for(int V=(1<<k)-1;V;V--){
for(int A=V;A;A=(A-1)&V){
f[u][V]=max(f[u][V],f[u][V-A]+f[v][A]);
if(V-A)ans=max(ans,1ll*f[u][V-A]+f[v][A]);
}
}
}
}
int main(){
read(n),read(k);
for(int i=1;i<=n;i++)read(a[i]);
for(int i=1;i<n;i++){
int u,v;ll w;
read(u),read(v),read(w);
ve[u].push_back(node(v,w));
ve[v].push_back(node(u,w));
}
for(t=1;t<=T;t++){
for(int i=1;i<=n;i++)c[i]=rand()%k;
for(int i=1;i<=n;i++)b[i]=c[a[i]];
dfs(1,0);
}
cout<<ans;
}