PROBLEM LINK:
Author: Yash Hiran
Editorialist: Yash Hiran
DIFFICULTY:
SIMPLE-EASY
PREREQUISITES:
dfs, trees
PROBLEM:
You are given a tree with n nodes, n-1 edges and a integer k. Suppose these k nodes are red coloured and remaining while coloured. You have to find the maximum sum of distances of k nodes from the root such that there are maximum number of white nodes in its shortest path to root.
QUICK EXPLANATION:
The contribution of each node will be the difference of its distance from root and its sub-tree size.
Maintain a set of these values. simply find the sum of last k elements from the set.
EXPLANATION:
The contribution of each leaf node will be its distance from the root. If we colour the non-leaf node red(i.e select in the list of k nodes), the numbers of whites nodes in the path of all the nodes in its sub-tree reduced by one.
Hence we can conclude that the contribution of each node is difference between its distance from root and sub-tree size. Create a array and store the above difference of each node in a array. Sort the array in descending order and find the sum of the first k elements.
SOLUTIONS:
#include<bits/stdc++.h>
using namespace std;
#define ll long long
#define maxN 100005
vector g[maxN];
bool vis[maxN];
vector delivery;
ll dfs(int node, int dis) {
vis[node] = true;
ll subTreeSize = 0;
for (auto x : g[node]) {
if (!vis[x]) {
subTreeSize += dfs(x , dis + 1);
}
}
delivery.push_back(dis - subTreeSize);
return subTreeSize + 1;
}
int main()
{
delivery.clear();
for (int i = 0; i < maxN; ++i) {
vis[i] = false;
g[i].clear();
}
ll n, k;
cin >> n >> k;
for (int i = 0; i < n - 1; ++i) {
int u, v; cin >> u >> v;
g[u].push_back(v);
g[v].push_back(u);
}
dfs(1, 0);
ll ans = 0;
sort(delivery.begin(), delivery.end(), greater<ll>());
int pos = 0;
while (k) {
ans += delivery[pos];
pos++;
k--;
}
cout << ans << endl;
return 0;
}