树的重心 举办聚会

 

#include<bits/stdc++.h>
using namespace std;
using ll=long long;
const int N=1e4+9;
int a[N],dp[N],sz[N],dep[N],n,k;
int sum=0;
vector<int> e[N] ;

void dfs1(int x,int fa)
{
  sz[x]=a[x];
  for(const auto & y : e[x])
  {
    if(y==fa)continue;
    dep[y]=dep[x]+1;
    dfs1(y,x);
    sz[x]+=sz[y];
  }
}

void dfs2(int x,int fa)
{
  for(const auto &y : e[x])
  {
    if(y==fa) continue;
    dp[y]=dp[x]-sz[y]+sum-sz[y];
    dfs2(y,x);
  }
}

int main()
{
  cin>>n>>k;
  for(int i=1;i<=n;i++) 
  {
    cin>>a[i];
    sum+=a[i];
  }

  for(int i=1;i<n;i++)
  {
    int u,v;cin>>u>>v;
    e[u].push_back(v);
    e[v].push_back(u);
  }
  dfs1(1,0);
  for(int i=1;i<=n;i++) dp[1]+=dep[i];
  dfs2(1,0);
  vector<int> v;
  for(int i=1;i<=n;i++) v.push_back(i);
  sort(v.begin(),v.end(),[](int &u,int &v)//不断的将两个数进行比较
  {
    if(dp[u]!=dp[v]) return dp[u]<dp[v];
    else return u<v;
    //对u和v进行排序,如果以u,v为下标在dp数组中的值不相等,则将u和v按照值小的那个排序
    //如果相等则按照u<v进行排序
    
  });
  cout<<v[k-1]<<endl;
  return 0;
}

---

#include<bits/stdc++.h>
using namespace std;
const int N=1e5+9;
long long n,k,dp[N],dep[N],sz[N],b[N],sum=0;
//sz[x]以x为根的所有人数之和 
//dp[j]以j为结点需要的路途总和
vector<int>a[N];

void dfs1(int x,int fa)
{
  sz[x]=b[x];//初始化人数 
  for(const auto & y : a[x] )
  {
    if(y==fa)continue;
    dep[y]=dep[x]+1;
    dfs1(y,x);
    sz[x]+=sz[y];//该节点及其子结点的所有人数 
  }
}

void dfs2(int x,int fa)
{
  for(const auto & y : a[x])
  {
    if(y==fa) continue;
    dp[y]=dp[x]-sz[y]+sum-sz[y];//该节点的所有人数均-1,其他人数距离+1 
    dfs2(y,x);
  }
}

int main()
{
  ios::sync_with_stdio(0),cin.tie(0),cout.tie(0);
  cin>>n>>k;
  for(int i=1;i<=n;i++)
  {
    cin>>b[i];//每个结点的权重
    sum+=b[i];
  }
  for(int k=1;k<n;k++)
  {
    int i,j;cin>>i>>j;
    a[i].push_back(j);
    a[j].push_back(i);
  }
  dfs1(1,0);
  for(int i=1;i<=n;i++) dp[1]+=dep[i];
  dfs2(1,0);
  vector<int> v;
  for(int i=1;i<=n;i++) v.push_back(i);
  sort(v.begin(),v.end(),[](int &u,int &v)
  {//自定义排序规则
    if(dp[u]!=dp[v]) return dp[u]<dp[v];
    else return u<v;
  });
  cout<<v[k-1]<<endl;//向量从0开始
  return 0;
}