BOJ 17532, ICPC Brazil Subregional 2019D Denouncing Mafia

#include <bits/stdc++.h>
#define usecppio ios::sync_with_stdio(0);cin.tie(0);cout.tie(0);
using namespace std;

int N, K;
vector <int> T[101010];
int depth[101010];
bool vst[101010];
struct path {
    int head;
    int length;
    int tail;
    bool operator<(const path &o) const {
        return length < o.length;
    }
};
priority_queue<path> pq;
path dfs(int r, int p) {
    depth[r] = depth[p] + 1;
    path longest_path;
    longest_path.length = 1;
    longest_path.head = r;
    longest_path.tail = r;
    for (int c : T[r]) {
        if (c == p) continue;
        path child_path = dfs(c, r);
        if (child_path.length + 1 > longest_path.length) {
            longest_path.length = child_path.length + 1;
            longest_path.tail = child_path.tail;
        }
    }
    pq.push(longest_path);
    return longest_path;
}

int parent[101010];
int32_t main() {
    usecppio
    cin >> N >> K;
    for (int i = 2; i <= N; i++) {
        int x; cin >> x;
        T[i].push_back(x);
        T[x].push_back(i);
        parent[i] = x;
    }
    dfs(1, 0);
    int ans = 0;
    while (!pq.empty() and K) {
        auto longest = pq.top();
        pq.pop();
        if (vst[longest.head]) continue;
        ans += longest.length;
        K -= 1;
        for (int cur = longest.tail; cur != longest.head; cur = parent[cur]) {
            vst[cur] = true;
        }
        vst[longest.head] = true;
    }
    cout << ans << '\n';
}