Hatena::Grouptopcoder

kuuso1の日記

2014-10-15

HackerRank Weekly Challenges - Week 11 Tree Pruning

| 00:28 | HackerRank Weekly Challenges - Week 11 Tree Pruning - kuuso1の日記 を含むブックマーク はてなブックマーク - HackerRank Weekly Challenges - Week 11 Tree Pruning - kuuso1の日記

https://www.hackerrank.com/contests/w11/challenges/tree-pruning

  • 各ノードに重みをもつ根付き木が与えられる(ノードは1~N, 1が根)
  • 任意のノードについてそのノード以下を取り除く操作が最大K回可能とする。
  • 残った木の重さを最大化する

(制約)2 ≤ N ≤ 10^5, 1 ≤ K ≤ 200, -10^9 ≤ 重み ≤ 10^9

方針:まず各ノードにぶら下がる重みを計算する。次に、どこを切れば重みが増えるかを考えると、

・あるノードを切ると、そこから下の重みが全てキャンセルされる

・切らなければ重みは変わらないので、そのノードにぶら下がる子ノードの切り方によって増やせる重みがそのまま伝わる

なので、自ノード以下でk(0 ≤ k ≤ K)回切った時に増やせる重みの最大値を覚えておいて、子から親の向きにdpする。

f:id:kuuso1:20141016002718p:image

コードは以下

using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;
 
class Solution{
	static void Main(){
		new Sol().Solve();
	}
}


class Sol{
	public void Solve(){
		
		Queue<int> Q=new Queue<int>();
		depth[root]=0;
		parent[root]=-1;
		Q.Enqueue(root);
		long minf=(long)-1e18;
		
		//深さと親を決める(Nが大きいのでdfsは回避)
		while(Q.Count>0){
			int x=Q.Dequeue();
			foreach(int l in E[x]){
				if(l!=parent[x]){
					depth[l]=depth[x]+1;
					parent[l]=x;
					Q.Enqueue(l);
				}
			}
		}
		
		int maxdepth=depth.Max();
		List<int>[] Dep=new List<int>[maxdepth+1];
		for(int i=0;i<=maxdepth;i++)Dep[i]=new List<int>();
		for(int i=0;i<N;i++){
			Dep[depth[i]].Add(i);
		}

		//各ノードの重みの合計を計算
		long[] sum=new long[N];
		for(int d=maxdepth;d>=0;d--){
			for(int j=0;j<Dep[d].Count;j++){
				sum[Dep[d][j]]=weight[Dep[d][j]];
				foreach(int l in E[Dep[d][j]]){
					if(l!=parent[Dep[d][j]]){
						sum[Dep[d][j]]+=sum[l];
					}
				}
			}
		}
		
		
		long[][] dp=new long[N][];
		for(int i=0;i<N;i++){
			dp[i]=new long[K+1];
			for(int j=0;j<=K;j++)dp[i][j]=minf;
		}
		
		//深い側から親の方に順に配るdp
		for(int d=maxdepth;d>=0;d--){
			for(int j=0;j<Dep[d].Count;j++){
				int now=Dep[d][j];
				dp[now][0]=0;
				foreach(int k in E[now]){
					if(k!=parent[now]){
						for(int ii=K;ii>=0;ii--){
							if(dp[now][ii]==minf)continue;
							for(int jj=K;jj>=0;jj--){
								if(dp[k][jj]==minf)continue;
								if(ii+jj>K)continue;
								dp[now][ii+jj]=Math.Max(dp[now][ii+jj],dp[now][ii]+dp[k][jj]);
							}
						}
					}
				}
				dp[now][1]=Math.Max(-sum[now],dp[now][1]);
			}
		}
		
		long Max=(long)-1e18;
		for(int i=0;i<=K;i++){
			Max=Math.Max(Max,dp[root][i]);
		}
		Console.WriteLine(sum[root]+Max);
		
	}
	
	int N;
	int K;
	int root;
	List<int>[] E;
	int[] parent;
	int[] depth;
	long[] weight;
	
	public Sol(){
		var d=ria();
		N=d[0];K=d[1];
		weight=rla();
		E=new List<int>[N];
		for(int i=0;i<N;i++)E[i]=new List<int>();
		parent=new int[N];
		root=0;
		depth=new int[N];
		for(int i=0;i<N-1;i++){
			var dd=ria();
			E[dd[0]-1].Add(dd[1]-1);
			E[dd[1]-1].Add(dd[0]-1);
		}
	}
	static int[] ria(){return Array.ConvertAll(Console.ReadLine().Split(' '),e=>int.Parse(e));}
	static long[] rla(){return Array.ConvertAll(Console.ReadLine().Split(' '),e=>long.Parse(e));}

}

editorialや慣れた人々はO(NK)で計算されてたのでO(NK^2)のコードを載せるのは少し気恥ずかしいですが、

まぁ初めて木dpを独力で通せたという事で!(TLE 3secに対して1.6secくらい)