LA 3902网络 无根树转有根树 树形dp - 放置最少的服务器距离k覆盖所有叶节点

#include <bits/stdc++.h>
#define LL long long
using namespace std;

int n, s, k;
vector<int> g[1005], node[1005];
int fa[1005];
bool vis[1005];

//无根树转有根树，计算fa数组， 根据深度把叶子节点插入nodes表中
void dfs(int u, int f, int d){
    fa[u]=f;
    int Len=g[u].size();
    if(Len==1&&d>k){    //叶子节点并且深度d>K
        node[d].push_back(u);//保存深度为d的叶子节点
    }
    for(int i=0; i<Len; i++){
        int v=g[u][i];
        if(v!=f){   //防止返回父节点
            dfs(v, u, d+1);
        }
    }
}

void dfs2(int u, int f, int d){//对一个服务器管理的节点染色

    vis[u]=1;//染色
    int Len=g[u].size();
    for(int i=0; i<Len; i++){
        int v=g[u][i];
        if(v!=f&&d<k){  //覆盖的条件
            dfs2(v, u, d+1);
        }
    }
}

int solve(){
    int ans=0;
    memset(vis, 0, sizeof(vis));
    for(int d=n-1; d>k; d--){//需要染色节点的条件
        for(int i=0; i<node[d].size(); i++){
            int u=node[d][i];
            if(vis[u]){//如果已经染色
                continue;
            }

            int v=u;
            for(int j=0; j<k; j++){//v是u的k级祖先
                v=fa[v];
            }
            dfs2(v, -1, 0);//放服务器 fa传入-1是因为上下距离，不是子树距离
            ans++;//统计
        }
    }

    return ans;
}

int main()
{
    int T;scanf("%d", &T);
    while(T--){
        scanf("%d%d%d", &n, &s, &k);
        for(int i=1; i<=n; i++){
            g[i].clear(), node[i].clear();
        }
        for(int i=0; i<n-1; i++){
            int u, v;
            scanf("%d%d", &u, &v);
            g[u].push_back(v);
            g[v].push_back(u);
        }
        dfs(s, -1, 0);
        printf("%d\n", solve());
    }

    return 0;
}
/*
1
14 12 2
13 14
14 12
11 14
12 2
2 1
2 3
3 10
3 4
4 9
4 5
5 6
5 7
5 8
*/