結果

問題 No.1718 Random Squirrel
コンテスト
ユーザー tarattata1
提出日時 2021-10-23 17:02:10
言語 C++14
(gcc 13.3.0 + boost 1.87.0)
結果
AC  
実行時間 88 ms / 2,000 ms
コード長 3,733 bytes
コンパイル時間 2,931 ms
コンパイル使用メモリ 160,524 KB
実行使用メモリ 43,648 KB
最終ジャッジ日時 2024-09-25 08:10:06
合計ジャッジ時間 5,903 ms
ジャッジサーバーID
（参考情報） 		judge5 / judge2
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 2
other AC * 31
権限があれば一括ダウンロードができます

ソースコード

diff #
raw source code 

#include <iostream>
#include <string>
#include <cstring>
#include <cstdlib>
#include <cmath>
#include <algorithm>
#include <vector>
#include <set>
#include <map>
#include <queue>
#include <stack>
#include <list>
#include <iterator>
#include <cassert>
#include <numeric>
#include <functional>
#include <ctime>
#include <bitset>
#pragma warning(disable:4996) 

#define ATCODER
#ifdef ATCODER
#include <atcoder/all>
#endif

typedef long long ll;
typedef unsigned long long ull;
#define LINF  9223300000000000000
#define LINF2 1223300000000000000
#define LINF3 1000000000000
#define INF 2140000000
//const long long MOD = 1000000007;
const long long MOD = 998244353;

using namespace std;
#ifdef ATCODER
using namespace atcoder;
#endif

//
//  全方位木
//    ST：各nodeに対する部分木に関する情報。単位元も定義する
//    op：(ST,ST)→ST という関数を定義。各部分木の情報を寄せ集めるときに使う
//    op2：ST→ST   子の部分木を合成した後、そこからcurrentの部分木の情報を計算。
//  出力
//    vector<ST> vans： 各nodeについて、そこから全方向への部分木の情報を寄せ集めたもの
//

vector<int> flag;

struct ST
{
    ST() {
        dist[0] = dist[1] = 0;       // unit item
    }
    int dist[2];   // 0:go & back   1:go
};

ST op(const ST a, const ST b)
{
    ST c;
    c.dist[0] = a.dist[0] + b.dist[0];
    c.dist[1] = min(a.dist[1] + b.dist[0], a.dist[0] + b.dist[1]);
    return c;
}

ST op2(const ST a, int f)
{
    ST b = a;
    if (b.dist[1]) {
        b.dist[0] += 2;
        b.dist[1] += 1;
    }
    else if( f ) {
        b.dist[0] = 2;
        b.dist[1] = 1;
    }
    return b;
}

void dfs0(const vector<vector<int>>& g, int par, int curr, vector<ST>& vs0)
{
    ST st;
    int found = 0;
    for (int i = 0; i < (int)g[curr].size(); i++) {
        int ne = g[curr][i];
        if (par == ne) continue;
        dfs0(g, curr, ne, vs0);
        st = op(st, vs0[ne]);
    }
    vs0[curr] = op2(st, flag[curr]);
    return;
}

void dfs1(const vector<vector<int>>& g, int par, int curr, const vector<ST>& vs0, const ST st1, vector<ST>& vans)
{
    vector<int> v;
    for (int i = 0; i < (int)g[curr].size(); i++) {
        int ne = g[curr][i];
        if (par == ne) continue;
        v.push_back(ne);
    }

    int num = (int)v.size();
    vector<ST> sv0(num + 1), sv1(num + 1);
    for (int i = 0; i < num; i++) {
        sv0[i + 1] = op(sv0[i], vs0[v[i]]);
    }
    for (int i = num-1; i >= 0; i--) {
        sv1[i] = op(sv1[i+1], vs0[v[i]]);
    }
    for (int i = 0; i < num; i++) {
        int ne = v[i];
        ST st2 = op(sv0[i], sv1[i + 1]);
        ST st3 = op(st1, st2);
        ST st = op2(st3, flag[curr]);
        dfs1(g, curr, ne, vs0, st, vans);
    }
    vans[curr] = op(st1, sv0[num]);
    return;
}

vector<ST> reroot(const vector<vector<int>>& g)
{
    int n = (int)g.size();
    vector<ST> vs0(n), vans(n);
    dfs0(g, -1, 0, vs0);
    ST st;
    dfs1(g, -1, 0, vs0, st, vans);
    return vans;
}

void solve()
{
    int n, K;
    scanf("%d%d", &n, &K);
    flag.resize(n);

    vector<vector<int>> g(n);  // to
    for (int i = 0; i < n - 1; i++) {
        int a, b;
        scanf("%d%d", &a, &b); a--; b--;
        g[a].push_back(b);
        g[b].push_back(a);
    }
    for (int i = 0; i < K; i++) {
        int t;
        scanf("%d", &t); t--;
        flag[t] = 1;
    }

    vector<ST> vans = reroot(g);

    for (int i = 0; i < n; i++) {
        printf("%d\n", vans[i].dist[1]);
    }

    return;
}

int main()
{
#if 1
    solve();
#else
    int T, t;
    scanf("%d", &T);
    for (t = 0; t < T; t++) {
        //printf("Case #%d: ", t + 1);
        solve();
    }
#endif
    return 0;
}
0