洛谷 P3258 [JLOI2014] 松鼠的新家

// P3258 [JLOI2014] 松鼠的新家 —— 正解：倍增求 LCA ＋ **点差分**，最后一趟子树求和
//
// ★★★ 这道题和 [P3128](/sol/p3128/) 只差一句话，而那一句话就是全部难点。
//   P3128：K 条路径，每条把沿途所有点 +1 —— 四次单点加就完了。
//   这道题：维尼**连着走** a₁ → a₂ → … → aₙ，而**接头处那个房间只经过一次**。
//
//   ⇒ 第 i 段（aᵢ → aᵢ₊₁）和第 i+1 段（aᵢ₊₁ → aᵢ₊₂）**都把 aᵢ₊₁ 算了一次**，
//     可维尼在那儿只停了一次 ⇒ **每个接头都要减掉一次**。
//   ⇒ 而最后那个房间 aₙ 是餐厅，题面明说「不需要再拿糖果吃」⇒ **它也要减一次**。
//
// ★★ 而这两件事能写成**同一句话**：对 i = 1 … n−1，`ans[a[i+1]]--`。
//   （i+1 ≤ n−1 时减的是「接头被算了两遍」，i+1 = n 时减的是「餐厅不拿糖」。）
//
// ⇒ 全部动作：n−1 条路径 × 四次单点加 ＋ n−1 次「接头 −1」＋ 一趟子树求和。O((n + n) log n)。

#include <bits/stdc++.h>
using namespace std;

const int MAXN = 300005;
const int LOG = 19;                       // 2^18 = 262144 < 3×10⁵ ≤ 2^19

int head_[MAXN], nxt_[MAXN * 2], to_[MAXN * 2], ecnt;
int up[MAXN][LOG], dep[MAXN], order_[MAXN], d[MAXN], a[MAXN];
int n;

inline void addEdge(int u, int v) { to_[++ecnt] = v; nxt_[ecnt] = head_[u]; head_[u] = ecnt; }

int lca(int u, int v) {
    if (dep[u] < dep[v]) swap(u, v);
    int diff = dep[u] - dep[v];
    for (int j = 0; j < LOG; j++) if ((diff >> j) & 1) u = up[u][j];
    if (u == v) return u;
    for (int j = LOG - 1; j >= 0; j--) if (up[u][j] != up[v][j]) { u = up[u][j]; v = up[v][j]; }
    return up[u][0];
}

int main() {
    if (scanf("%d", &n) != 1) return 0;
    for (int i = 1; i <= n; i++) if (scanf("%d", &a[i]) != 1) return 0;
    for (int i = 0; i < n - 1; i++) {
        int x, y;
        if (scanf("%d %d", &x, &y) != 2) return 0;
        addEdge(x, y); addEdge(y, x);     // ⚠ 无向边，两边都要存
    }

    /* BFS 求 dep 和 up[·][0]（⚠ 别用递归：顶格 3×10⁵ 排成一条链就是 30 万层）*/
    {
        vector<char> vis(n + 1, 0);
        int cnt = 0;
        order_[cnt++] = 1; vis[1] = 1; dep[1] = 1;
        up[1][0] = 0;                     // ⚠⚠ 根的父亲是 0 号点，不是它自己
        for (int i = 0; i < cnt; i++) {
            int u = order_[i];
            for (int e = head_[u]; e; e = nxt_[e]) {
                int v = to_[e];
                if (vis[v]) continue;
                vis[v] = 1; dep[v] = dep[u] + 1; up[v][0] = u; order_[cnt++] = v;
            }
        }
    }
    for (int j = 1; j < LOG; j++)
        for (int v = 1; v <= n; v++)
            up[v][j] = up[up[v][j - 1]][j - 1];

    for (int i = 1; i < n; i++) {
        int s = a[i], t = a[i + 1];
        int l = lca(s, t);
        d[s]++; d[t]++; d[l]--;
        if (up[l][0]) d[up[l][0]]--;      // ★ l 是根时 up[l][0] == 0，这笔账没人要收
    }

    /* 自底向上把子树和累起来 —— ⚠ 必须**逆着** BFS 序走 */
    for (int i = n - 1; i >= 1; i--) {
        int u = order_[i];
        d[up[u][0]] += d[u];
    }
    /* ★★★ 接头（以及最后那个餐厅）只算一次 —— ⚠⚠ 这一刀必须落在**还原之后**：
       差分数组里 d[u] 管的是「u 的整棵子树」，`d[t]--` 会把 t 头上的每一个祖先也一起减掉。*/
    for (int i = 1; i < n; i++) d[a[i + 1]]--;

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

// P3258 的两笔账：这份数据长什么样，以及两条路各做多少次基本动作（机器无关）
//
// 用法：./p3258Count <csv|table> < 一份输入
//
// ★ 顺手数三个决定成败的东西：
//   ① **a₁ 和 aₙ 是不是同一个房间** —— 那正好是错法③（减在起点上）的触发条件；
//   ② **有几个接头不是根** —— 那是错法④（把「接头 −1」写进差分数组）的触发条件；
//   ③ **暴力一共要走进多少个房间** —— 顶格那张表的主角。
#include <bits/stdc++.h>
using namespace std;

const int MAXN = 300005;
const int LOG = 19;
static int head_[MAXN], nxt_[MAXN * 2], to_[MAXN * 2], ecnt;
static int up[MAXN][LOG], dep[MAXN], order_[MAXN], d[MAXN], a[MAXN];
static int n;

static int lca(int u, int v) {
    if (dep[u] < dep[v]) swap(u, v);
    int diff = dep[u] - dep[v];
    for (int j = 0; j < LOG; j++) if ((diff >> j) & 1) u = up[u][j];
    if (u == v) return u;
    for (int j = LOG - 1; j >= 0; j--) if (up[u][j] != up[v][j]) { u = up[u][j]; v = up[v][j]; }
    return up[u][0];
}

int main(int argc, char** argv) {
    string mode = argc > 1 ? argv[1] : "table";
    if (scanf("%d", &n) != 1) return 0;
    for (int i = 1; i <= n; i++) if (scanf("%d", &a[i]) != 1) return 0;
    for (int i = 0; i < n - 1; i++) {
        int x, y;
        if (scanf("%d %d", &x, &y) != 2) return 0;
        to_[++ecnt] = y; nxt_[ecnt] = head_[x]; head_[x] = ecnt;
        to_[++ecnt] = x; nxt_[ecnt] = head_[y]; head_[y] = ecnt;
    }
    {
        vector<char> vis(n + 1, 0);
        int c = 0;
        order_[c++] = 1; vis[1] = 1; dep[1] = 1; up[1][0] = 0;
        for (int i = 0; i < c; i++) {
            int u = order_[i];
            for (int e = head_[u]; e; e = nxt_[e]) {
                int v = to_[e];
                if (vis[v]) continue;
                vis[v] = 1; dep[v] = dep[u] + 1; up[v][0] = u; order_[c++] = v;
            }
        }
    }
    for (int j = 1; j < LOG; j++) for (int v = 1; v <= n; v++) up[v][j] = up[up[v][j - 1]][j - 1];

    int maxDep = 0;
    for (int i = 1; i <= n; i++) maxDep = max(maxDep, dep[i]);

    long long bruteSteps = 1;                 // 出发那一间也算「走进去」一次
    long long jointNotRoot = 0, lcaRoot = 0;
    for (int i = 1; i < n; i++) {
        int s = a[i], t = a[i + 1];
        int l = lca(s, t);
        if (l == 1) lcaRoot++;
        if (t != 1) jointNotRoot++;
        bruteSteps += (long long)(dep[s] - dep[l]) + (dep[t] - dep[l]);   // ⚠ 起点不重复计
        d[s]++; d[t]++; d[l]--;
        if (up[l][0]) d[up[l][0]]--;
    }
    for (int i = n - 1; i >= 1; i--) d[up[order_[i]][0]] += d[order_[i]];
    for (int i = 1; i < n; i++) d[a[i + 1]]--;

    int mx = 0, mn = 1 << 30;
    long long tot = 0;
    for (int i = 1; i <= n; i++) { mx = max(mx, d[i]); mn = min(mn, d[i]); tot += d[i]; }

    /* 差分那条路的基本动作：建表 n·(LOG−1) 次写 ＋ 每段最多 2·LOG 次读 ＋ 4(n−1) 次单点加 */
    long long fastSteps = (long long)n * (LOG - 1) + (long long)(n - 1) * (2 * LOG + 4);

    vector<pair<string, string> > out;
    out.push_back(make_pair("n", to_string(n)));
    out.push_back(make_pair("maxdep", to_string(maxDep)));
    out.push_back(make_pair("head_eq_tail", to_string(a[1] == a[n] ? 1 : 0)));
    out.push_back(make_pair("joint_not_root", to_string(jointNotRoot)));
    out.push_back(make_pair("lca_root", to_string(lcaRoot)));
    out.push_back(make_pair("ans_max", to_string(mx)));
    out.push_back(make_pair("ans_min", to_string(mn)));
    out.push_back(make_pair("ans_sum", to_string(tot)));
    out.push_back(make_pair("brute_steps", to_string(bruteSteps)));
    out.push_back(make_pair("fast_steps", to_string(fastSteps)));
    out.push_back(make_pair("ratio", to_string(bruteSteps / fastSteps)));
    if (mode == "csv") for (size_t i = 0; i < out.size(); i++) printf("%s,%s\n", out[i].first.c_str(), out[i].second.c_str());
    else for (size_t i = 0; i < out.size(); i++) printf("  %-16s %s\n", out[i].first.c_str(), out[i].second.c_str());
    return 0;
}