E. Paint the Tree
题意:
输入
n
,
k
(
5
e
5
)
n,k(5e5)
n,k(5e5);
接下来
n
−
1
n-1
n−1行,每行
u
,
v
,
w
u,v,w
u,v,w,表示树边;
每个点有
k
k
k个颜色,每种颜色出现次数不超过两次,当一条边两端有颜色相同,这条边的权值算作贡献,问贡献和最大为多少?
题解:
d
p
[
u
]
[
0
]
dp[u][0]
dp[u][0]表示
u
u
u子树不和父亲节点有相同颜色的贡献和最大值;
d
p
[
u
]
[
1
]
dp[u][1]
dp[u][1]表示
u
u
u子树和父亲节点有相同颜色的贡献最大值;
那么如何转移?
对于节点
u
u
u,有
v
1
,
v
2
,
…
,
v
x
v_1,v_2,\dots,v_x
v1,v2,…,vx儿子
贪心的选择比较优的儿子就好,注意用优先队列维护。
代码:
#include<bits/stdc++.h>
using namespace std;
const int N=5e5+7;
#define ll long long
int q,n,k;
struct Edge{
int v,w,nxt;
}e[N<<1];
int head[N],cnt;
inline void add(int u,int v,int w){
e[cnt]=(Edge){v,w,head[u]};
head[u]=cnt++;
}
ll dp[N][2];
void dfs(int u,int fa){
priority_queue<ll>q;ll sum=0;
for(int i=head[u];~i;i=e[i].nxt){
Edge e1=e[i];
if(e1.v==fa)continue;
dfs(e1.v,u);
ll t=dp[e1.v][1]+e1.w-dp[e1.v][0];
if(t>0)q.push(t);
sum+=dp[e1.v][0];
}
if(k-1==0)dp[u][1]=sum;
int cnt=0;ll he=0;
while(!q.empty()){
he+=q.top();q.pop();cnt++;
if(cnt==k-1)dp[u][1]=sum+he;
if(cnt==k)dp[u][0]=sum+he;
}
if(cnt<k-1)dp[u][1]=sum+he;
if(cnt<k)dp[u][0]=sum+he;
}
void solve(){
for(int i=1;i<=n;i++)dp[i][0]=dp[i][1]=0;
dfs(1,0);
printf("%lld\n",max(dp[1][0],dp[1][1]));
}
int main(){
// freopen("tt.in","r",stdin),freopen("tt.out","w",stdout);
scanf("%d",&q);
while(q--){
scanf("%d%d",&n,&k);
for(int i=1;i<=n;i++)head[i]=-1;
// memset(head,-1,sizeof(head));就是因为这个超时了。
cnt=0;
for(int i=1;i<n;i++){
int u,v,w;
scanf("%d%d%d",&u,&v,&w);
add(u,v,w),add(v,u,w);
}
solve();
}
return 0;
}