CS Academy City Attractions

リンク : https://csacademy.com/contest/archive/task/city-attractions/statement/

問題概要

N頂点の木が存在する。それぞれの頂点に美しさA[i]として入力される。d(x,y):=頂点xと頂点yの距離とする。頂点xにいるとき、次に進む場所はA[y]-d(x,y)が最大になる都市である。K回進んだとき、最後の頂点はいくつか。

解説

すぐに全方位木dpっぽいなとは思ったが、MLEがつらすぎた... 全方位木dpとは簡単に説明すると、全頂点を根として一回探索したいときにO(N²)ではなく、O(N)でできる考え方である。一度全方位木dpをして、頂点iにいるときに次に進む場所を計算しておく。これができれば、あとは頂点1から順番に見ていってサイクルが必ず存在するためKをサイクルの大きさで割った余りにして、また前から見ていけばO(N)で答えを求めることができる。

LL n, k, a[300010], h[300010];
vector<LL>g[300010];
Pll nex[300010], dp[300010];
bool f[300010];
void cmp(Pll &x,Pll y) {
    if (x.fir < y.fir||(x.fir==y.fir&&x.sec>y.sec))
        x = y;
}
void dfs1(LL cur, LL par) {
    nex[cur] = Pll{ -INF / 2,-1 };
    rep(i, g[cur].size()) {
        LL v = g[cur][i];
        if (v == par)continue;
        dfs1(v, cur);
        cmp(nex[cur], { nex[v].fir - 1,nex[v].sec });
        cmp(nex[cur], { a[v] - 1,v });
    }
}
void dfs2(LL cur, LL par) {
    if (par != -1) {
        cmp(dp[cur], { dp[par].fir - 1, dp[par].sec});
        cmp(dp[cur], Pll{ a[par] - 1,par });
    }
    Pll mx = Pll{ -INF / 2,-1 };
    rep(i, g[cur].size()) {
        LL v = g[cur][i];
        if (v == par)continue;
        cmp(dp[v], { mx.fir - 2,mx.sec });
        cmp(mx, nex[v]);
        cmp(mx, { a[v],v });
    }
    reverse(g[cur].begin(), g[cur].end());
    mx = Pll{ -INF / 2,-1 };
    rep(i, g[cur].size()) {
        LL v = g[cur][i];
        if (v == par)continue;
        cmp(dp[v], { mx.fir - 2,mx.sec});
        cmp(mx, nex[v]);
        cmp(mx, { a[v],v });
    }
    rep(i, g[cur].size()) {
        LL v = g[cur][i];
        if (par == v)continue;
        dfs2(v, cur);
    }
    cmp(dp[cur], nex[cur]);
    h[cur] = dp[cur].sec;
}
int main() {
    cin >> n>>k;
    rep(i, n) {
        cin >> a[i];
    }
    rep(i, n - 1) {
        LL x, y;
        cin >> x >> y;
        x--, y--;
        g[x].push_back(y);
        g[y].push_back(x);
    }
    rep(i, n)dp[i] = Pll{ -INF / 2,-1 };
    dfs1(0, - 1);
    dfs2(0, -1);
    LL to = 0;
    while (f[to]==0) {
        f[to] = 1;
        to = h[to];
        k--;
        if (k == 0) {
            cout << to + 1 << endl;
            return 0;
        }
    }
    rep(i, n)f[i] = 0;
    LL cnt = 0;
    while (f[to] == 0) {
        f[to] = 1;
        to = h[to];
        cnt++;
    }
    k %= cnt;
    while (k > 0)to = h[to], k--;
    cout << to + 1 << endl;
    return 0;
}