字节跳动夏令营2021笔试题t3

一个公司为每一个员工买衣服，有D,E,F三种衣服，每一种衣服给员工i的开心值分别为Di，Ei，Fi，但是员工不愿意和他的直系boss衣服一样，求所有人开心值和的最大值

输入：
n，代表n个员工
n行，第i行 为Di，Ei，Fi
下面n-1行 每一行为 a b 代表a是b的直属boss

样例输入：
3
2 4 9
1 3 5
1 2 3
0 1
0 2
样例输出：
14

0号选F，1号选E,2号选E，9+3+2=14

---

树状dp

• 员工和直属领导的关系构成了一棵树，并且我们使用0,1,2表示衣服D E F
• 我们使用dp[i][j]表示 编号为i的人选择衣服j时以i为根节点的子树开心值之和
• 使用深度优先搜索，每一个根节点的dp值等于所有子节点的最大值之和
  即

dp[i][j] = dp[i][j] + max(treeDP(nodes[i].nex,(j+1)%3),treeDP(nodes[i].nex,(j+2)%3);

• 该结点i叶结点，那么直接返回dp[i][j]

代码

#include<iostream>
#include<stdio.h>
#include<cstring>
#include<vector>
using namespace std;

int dp[5005][3];
int h[5005][3];
int f[5005];
int n;
int MAX(int a,int b){
    return a>b?a:b;
}

class node{
    public:
        int id;
        vector<int>nex;
}nodes[5005];

int treeDP(int tmp,int j){//编号i为根的子树 选的j号衣服 的开心值的最大值
    if(nodes[tmp].nex.size()==0){//如果是叶结点 直接返回
        // cout<<"leaf "<<tmp<<" "<<j<<"  "<<dp[tmp][j]<<endl;
        return dp[tmp][j];
    }
    // cout<<"In subTree "<<tmp<<" "<<j<<endl;
    for(int i=0;i<nodes[tmp].nex.size();i++){//加上每个叶结点的最大值
        dp[tmp][j]=dp[tmp][j]+(MAX(treeDP(nodes[tmp].nex[i],(j+1)%3),treeDP(nodes[tmp].nex[i],(j+2)%3)));// j+1 j+2 是为了确保选的衣服和j不一样
    }
    // cout<<"dp "<<tmp<<' '<<j<<"  = "<<dp[tmp][j]<<endl<<endl;
    return dp[tmp][j];
}

int getf(int x){//找到根节点用的
    if(f[x]==x)return x;
    else{
        return f[x]=getf(f[x]);
    }
}
void ini(){//初始化dp为输入的快乐值
    for(int i=0;i<n;i++){
        dp[i][0]=h[i][0];
        dp[i][1]=h[i][1];
        dp[i][2]=h[i][2];
    }
}
int main(){
    
    int root;
    cin>>n;
    for(int i=0;i<n;i++){
        cin>>h[i][0]>>h[i][1]>>h[i][2];
    }

    for(int i=0;i<n-1;i++){
        int father,son;
        cin>>father>>son;
        nodes[father].nex.push_back(son);//构造树
        f[son]=father;
        if(i==n-2){
            root = getf(father);
        }
    }
    // cout<<"root "<<root<<endl;
    int ans = 0;
    ini();
    ans = MAX(ans,treeDP(root,0));
    ini();
    ans = MAX(ans,treeDP(root,1));
    ini();
    ans = MAX(ans,treeDP(root,2));
    cout<<ans<<endl;
    system("pause");
    return 0;
}