LCA ライブラリを持っていれば書くだけ！

問題概要

$N$ 頂点の根付き木が与えられる (根は頂点 0)。この木上で幅優先探索を行う (隣接する頂点のうち、頂点番号が小さい順にキューに push する)。

このとき、探索順序が $v_{1} (= 0), v_{2}, \dots, v_{N}$ であったとしたとき、

$\sum_{i = 1}^{N-1} (頂点 v_{i} と頂点 v_{i+1} の距離)$

の値を求めよ。

制約

$1 \le N \le 10^{5}$

考えたこと

木の 2 頂点間の距離を求める問題は、前処理を行なっておけば $O(\log N)$ で実現できることが知られている。

ライブラリで殴る。計算量は $O(N \log N)$ 。

#include <bits/stdc++.h>
using namespace std;
template<class T> inline bool chmax(T& a, T b) { if (a < b) { a = b; return 1; } return 0; }
template<class T> inline bool chmin(T& a, T b) { if (a > b) { a = b; return 1; } return 0; }

using Graph = vector<vector<int>>;
struct RunTree {
    // id[v][w] := the index of node w in G[v]
    vector<map<int,int> > id;

    // num[v][i] := the size of subtree of G[v][i] with parent v
    vector<vector<long long> > num;

    // constructor
    RunTree(const Graph &G, int root = 0) {
        init(G, root);
    }

    // size(u, v) := the size of subtree v with parent u
    long long size(int u, int v) {
        return num[u][id[u][v]];
    }

    // lca(u, v)
    int getLCA(int u, int v) {
        if (depth[u] > depth[v]) swap(u, v);
        for (int i = 0; i < (int)parent.size(); ++i)
            if ( (depth[v] - depth[u]) & (1<<i) )
                v = parent[i][v];
        if (u == v) return u;
        for (int i = (int)parent.size()-1; i >= 0; --i) {
            if (parent[i][u] != parent[i][v]) {
                u = parent[i][u];
                v = parent[i][v];
            }
        }
        return parent[0][u];
    }

    // length(u, v)
    long long length(int u, int v) {
        int lca = getLCA(u, v);
        return depth[u] + depth[v] - depth[lca]*2;
    }

    // getParent(v, p) := the parent of v directed for p
    int getParent(int v, int p) {
        if (v == p) return -1;
        int lca = getLCA(v, p);
        if (lca != v) return parent[0][v];
        for (int i = (int)parent.size()-1; i >= 0; --i) {
            if (parent[i][p] != -1 && depth[parent[i][p]] > depth[v]) {
                p = parent[i][p];
            }
        }
        return p;
    }
    
    // rec
    vector<vector<int> > parent;
    vector<int> depth;
    int rec(const Graph &G, int v, int p = -1, int d = 0) {
        int p_index = -1;
        int sum = 1;
        parent[0][v] = p;
        depth[v] = d;
        for (int i = 0; i < (int)G[v].size(); ++i) {
            int ch = G[v][i];
            id[v][ch] = i;
            if (ch == p) {
                p_index = i;
                continue;
            }
            int s = rec(G, ch, v, d+1);
            num[v][i] = s;
            sum += s;
        }
        if (p_index != -1) num[v][p_index] = (int)G.size() - sum;
        return sum;
    }

    // init
    void init(const Graph &G, int root = 0) {
        int N = (int)G.size();
        id.assign(N, map<int,int>());
        num.assign(N, vector<long long>());
        for (int v = 0; v < N; ++v) num[v].assign((int)G[v].size(), 0);
        int V = (int)G.size();
        int h = 1;
        while ((1<<h) < N) ++h;
        parent.assign(h, vector<int>(N, -1));
        depth.assign(N, -1);
        rec(G, root);
        for (int i = 0; i+1 < (int)parent.size(); ++i)
            for (int v = 0; v < V; ++v)
                if (parent[i][v] != -1)
                    parent[i+1][v] = parent[i][parent[i][v]];
    }
};

int main() {
    int N;
    cin >> N;
    Graph G(N);
    for (int i = 1; i < N; ++i) {
        int p; cin >> p; --p;
        G[i].push_back(p);
        G[p].push_back(i);
    }
    RunTree rt(G);
    long long res = 0;
    vector<long long> dist(N, -1);
    queue<int> que;
    dist[0] = 0, que.push(0);
    int pv = -1;
    while (!que.empty()) {
        int v = que.front(); que.pop();
        if (pv != -1) res += rt.length(v, pv);
        pv = v;
        for (auto to : G[v]) {
            if (dist[to] == -1) {
                dist[to] = dist[v] + 1;
                que.push(to);
            }
        }
    }
    cout << res << endl;
}