OJ贪心#257. 树的颜色

题目描述

​ 有一棵树，它的所有节点都需要染色，每个节点都有一个代价基础值 Ci

。第一个染色的是根节点，其余节点染色的时候其父节点必须已染色。染色一个节点会用掉一个时间单位，每个节点染色的代价是染完此节点时的总时间 T 乘上这个节点的基础值 Ci

。求染完所有节点所需的最小代价。

---

输入

​ 第一行包含两个整数 N,R

其中，N 是树中节点个数，R

是根节点编号。

​ 第二行输入 N

个整数，编号为 i 的节点的代价基础值 Ci。（1≤Ci≤500）

​ 接下来 N−1

行为边的信息，每行两个数分别表示父节点编号和子节点编号。

输出

​ 输出一个整数，表示最小代价。

---

样例输入

5 1
1 2 1 2 4
1 2
1 3
2 4
3 5

样例输出

33

---

数据规模与约定

​ 时间限制：1 s

​ 内存限制：256 M

​ 100% 的数据保证 1≤R≤N≤1000

import java.util.*;

public class oj_treeColor {
    public static void main(String[] args){
        Scanner scan=new Scanner(System.in);
        int n=scan.nextInt(),r=scan.nextInt();
        int[] weight=new int[n];
        for(int i=0;i<n;i++){
            weight[i]=scan.nextInt();
        }
        int[][] relation=new int[n-1][];
        for(int i=0;i<n-1;i++){
            relation[i]=new int[2];
            for(int j=0;j<2;j++){
                relation[i][j]=scan.nextInt();
            }
        }
        MiniWeightSum miniWeightSum=new MiniWeightSum(weight,relation);
        System.out.println(miniWeightSum.outc);
        scan.close();
    }
}
class MiniWeightSum{
    int[] weight;
    int[][] relation;
    Map<Integer,Node> hashMap;
    Boolean[] visited;
    int outc=0,cnt;
    MiniWeightSum(int[] weight,int[][] relation){
        this.relation=relation;
        this.weight=weight;
        solve();
    }
    void solve(){
        hashMap=new HashMap<>();
        hashMap.put(1,new Node(weight[0],1,0));
        for (int i=0;i<relation.length;i++){    //hashmap初始化
            int seq=relation[i][1];
            int p=relation[i][0];
            int w=weight[seq-1];
            hashMap.put(seq,new Node(w,seq,p));
            Node pn=hashMap.get(p);
            pn.setChild(seq+"");
            sumofChildWeight(p,w);//计算每棵子树总权值
        }
        visited=new Boolean[relation.length+2];
        Arrays.fill(visited,false);
        cnt=1;
        outc=0;
        countMiniSum(1);
    }
    void countMiniSum(int root){
        visited[root]=true;
        Node now=hashMap.get(root);
        String s=now.child;
        outc+=cnt*hashMap.get(root).weight;
        cnt++;
        if(s=="")   return;
        for(int j=0;j<s.length();j++){
            int tag=0,index=0;
            for(int i=0;i<s.length();i++){
                int cntWeight=hashMap.get(s.charAt(i)-'0').cntWeight;
                if(!visited[s.charAt(i)-'0']&&cntWeight>tag){
                    tag=cntWeight;
                    index=s.charAt(i)-'0';
                }
            }
            countMiniSum(index);
        }
    }
    void sumofChildWeight(int current,int weight){
        if(current==0) return;
        Node now=hashMap.get(current);
        now.addWeight(weight);
        int parent=now.parent;
        sumofChildWeight(parent,weight);
    }
    static class Node{
        int weight,seq,parent,cntWeight;
        String child;
        Node(int weight,int seq,int parent){
            this.weight=weight;
            this.seq=seq;
            this.parent=parent;
            cntWeight=weight;
            child="";
        }
        void setChild(String s){
            child+=s;
        }
        void addWeight(int n){
            cntWeight+=n;
        }
    }
}