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);
}
}