PTREE - Cây P đỉnh ( Cơ bản )

Tác giả: khuc_tuan

Ngôn ngữ: Java

import java.io.*;
import java.util.*;

public class Main {
	static int n,p;
	static int[] c, father;
	static ArrayList<Integer>[] ke;
	static int[][] f,d;
	static void dfs(int i) {
		int socon = 0;
		for(int j:ke[i]) if(j!=father[i]) {
			father[j]=i;
			dfs(j);
			++socon;
		}
		if(socon==0) {
			f[i][1] = c[i];
			f[i][0] = 0;
			return;
		}
		int s=0;
		for(int j=1;j<=p;++j) d[0][j] = -20000000;
		for(int j:ke[i]) if(j!=father[i]) {
			++s;
			for(int t=0;t<=p;++t) {
				d[s][t]=-20000000;
				for(int k=0;k<=t;++k) d[s][t]=Math.max(d[s][t],d[s-1][k]+f[j][t-k]);
			}
		}
		for(int t=1;t<=p;++t) f[i][t] = d[socon][t-1]+c[i];
		f[i][0] = 0;
	}
	static void truyvet(int i,int p) {
		if(p==0) return;
		int socon = 0;
		ArrayList<Integer> con=new ArrayList();
		for(int j:ke[i]) if(j!=father[i]) {
			++socon;
			con.add(j);
		}
		if(socon==0) {
			System.out.print(i+" ");
			return;
		}
		System.out.print(i+" ");
		int s=0;
		for(int j=1;j<=p;++j) d[0][j] = -20000000;
		for(int j:ke[i]) if(j!=father[i]) {
			++s;
			for(int t=0;t<=p;++t) {
				d[s][t]=-20000000;
				for(int k=0;k<=t;++k) d[s][t]=Math.max(d[s][t],d[s-1][k]+f[j][t-k]);
			}
		}
		int x=p-1;
		Stack<Integer> ss = new Stack();
		for(int j=socon;j>=1;--j) {
			for(int k=0;k<=x;++k) if(d[j-1][k]+f[con.get(j-1)][x-k] == d[j][x]) {
				ss.push(con.get(j-1));
				ss.push(x-k);
				x = k;
				break;
			}
		}
		while(!ss.isEmpty()){
			int b=ss.pop(), a=ss.pop();
			truyvet( a, b);
		}
	}
	public static void main(String[] Args) throws Exception {
		BufferedReader kb=new BufferedReader(new InputStreamReader(System.in));
		StringTokenizer st=new StringTokenizer(kb.readLine());
		n = Integer.parseInt(st.nextToken());
		p = Integer.parseInt(st.nextToken());
		c = new int[n+1];
		st=new StringTokenizer(kb.readLine());
		for(int i=1;i<=n;++i) c[i] = Integer.parseInt(st.nextToken());
		ke = new ArrayList[n+1];
		for(int i=1;i<=n;++i) ke[i]=new ArrayList();
		for(int i=1;i<n;++i) {
			st=new StringTokenizer(kb.readLine());
			int u=Integer.parseInt(st.nextToken()), v=Integer.parseInt(st.nextToken());
			ke[u].add(v);
			ke[v].add(u);
		}
		f=new int[n+1][p+1];
		d=new int[n+1][p+1];
		for(int i=0;i<=n;++i) for(int j=0;j<=p;++j) f[i][j]=-20000000;
		father=new int[n+1];
		dfs(1);
		int res=0, imax=0;
		for(int i=1;i<=n;++i) if(res<f[i][p]) {
			res=f[i][p];
			imax=i;
		}
		truyvet(imax,p);
	}
}

Download