CSP 202312-3 树上搜索

#include <bits/stdc++.h>
using namespace std;
const int N = 2050;
using ll = long long;
struct node {
    int id;
    int fa;
    ll value;
    vector<int> child;
} nodes[N];
int n, m;
ll calc(int x, vector<ll>& cval, vector<bool>& vis) {
    ll sum = 0;
    for (auto c : nodes[x].child) {
        sum += calc(c, cval, vis);
    }

    if (!vis[x]) cval[x] = sum + nodes[x].value;
    return cval[x];
}
bool belong(int u, int x) {
    if (u == x) return true;
    for (auto& c : nodes[u].child) {
        if (c == x) return true;
        if (belong(c, x)) return true;
    }
    return false;
}
void solve(int x) {
    int cnt = n;
    vector<int> st;
    vector<bool> vis(n + 5, false);
    while (cnt > 1) {
        vector<ll> cval(n + 5, 0);
        vector<ll> w(n + 5, 0);
        calc(1, cval, vis);
        ll tot_cnt = 0;
        for (int i = 1; i <= n; i++) {
            if (!vis[i]) tot_cnt += nodes[i].value;
        }
        for (int i = 1; i <= n; i++) {
            if (!vis[i]) {
                w[i] = abs(tot_cnt - cval[i] - cval[i]);
            }
        }
        ll mn = 1e10;
        ll idx = -1;
        for (int i = 1; i <= n; i++) {
            if (!vis[i]) {
                if (w[i] < mn) {
                    mn = w[i];
                    idx = i;
                } else if (w[i] == mn) {
                    if (idx > i) idx = i;
                }
            }
        }
        st.push_back(idx);
        if (belong(idx, x)) {
            for (int i = 1; i <= n; i++) {
                if (!vis[i] && !belong(idx, i)) {
                    vis[i] = true;
                    cnt--;
                }
            }
        } else {
            for (int i = 1; i <= n; i++) {
                if (i == idx && !vis[idx]) {
                    vis[idx] = true;
                    cnt--;
                    continue;
                }
                if (!vis[i] && belong(idx, i)) {
                    vis[i] = true;
                    cnt--;
                }
            }
        }
    }
    for (auto c : st) {
        cout << c << ' ';
    }
    cout << '\n';
}
int main() {
    cin >> n >> m;
    for (int i = 1; i <= n; i++) {
        cin >> nodes[i].value;
    }
    for (int i = 2; i <= n; i++) {
        int fa;
        cin >> fa;
        nodes[i].fa = fa;
        nodes[fa].child.push_back(i);
    }
    nodes[1].fa = -1;
    for (int i = 1; i <= m; i++) {
        int x;
        cin >> x;
        solve(x);
    }
}