10544번: KAMP - 소스 최적화 필요

https://www.acmicpc.net/problem/10544

#include<stdio.h>
#include<vector>
#include<algorithm>
using namespace std;
const int MAX_N = 5e5;
typedef long long ll;
vector<pair<int, int> > adj[MAX_N + 1];
vector<int> vis;
int n, k, ind[MAX_N + 1];
ll dis[2][MAX_N + 1], sdis[MAX_N + 1];
bool ck[MAX_N + 1], sck[MAX_N + 1], tck[MAX_N + 1];
ll res;
void dfs1(int h) {
    if (tck[h] || ind[h] != 1) return;
    vis.push_back(h);
    sck[h] = true;
    for (auto it : adj[h])
        if (!sck[it.first]) --ind[it.first], dfs1(it.first);
}
ll tdis;
int t;
ll dfs2(int h, ll d) {
    ck[h] = true;
    if (tdis < d) tdis = d, t = h;
    for (auto it : adj[h])
        if (!ck[it.first] && !sck[it.first]) dfs2(it.first, d + it.second);
    ck[h] = false;
}
ll dfs3(int h, ll d, int st) {
    ck[h] = true;
    dis[st][h] = d;
    for (auto it : adj[h])
        if (!ck[it.first]) dfs3(it.first, d + it.second, st);
    ck[h] = false;
}
int main() {
    scanf("%d %d", &n, &k);
    for (int i = 0, a, b, c; i < n - 1; i++) {
        scanf("%d %d %d", &a, &b, &c);
        adj[a].push_back({ b,c });
        adj[b].push_back({ a,c });
        ind[a]++;
        ind[b]++;
        res += c * 2;
    }
    int x;
    for (int i = 0; i < k; i++) {
        scanf("%d", &x);
        tck[x] = true;
    }
    for (int i = 1; i <= n; i++)
        if (ind[i] == 1) dfs1(i);
    reverse(vis.begin(), vis.end());
    for (auto it : vis)
        for (auto h : adj[it])
            if (!sck[h.first] || sdis[h.first]) sdis[it] = sdis[h.first] + h.second, res -= 2 * h.second;
    int t1, t2;
    dfs2(x, 0); t1 = t;
    tdis = 0;
    dfs2(t1, 0); t2 = t;
    dfs3(t1, 0, 0);
    dfs3(t2, 0, 1);
    for (int i = 1; i <= n; i++) printf("%lld\n", res - max(dis[0][i], dis[1][i]) + 2 * sdis[i]);
    return 0;
}