[CF587C] Duff in the Army

[CF587C] Duff in the Army - 树上倍增,STL

Description

有 n 个城市，由 n-1 条边连接。两个城市之间的道路是唯一的。有 m 个人，住在这 n 个城市中，现在给出 m 个人生活的城市编号。你需要回答，从一个城市 u 到另一个城市 v 的路径中，编号前 a 小的人的编号是哪些。

Solution

树上倍增处理路径问题，重载结点加法

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

#define int long long

const int N = 200005;

int n, m, q;
vector<int> g[N];

struct Node
{
    vector<int> a;
    void insert(int x)
    {
        a.push_back(x);
        if (a.size() > 10)
            a.erase(max_element(a.begin(), a.end()));
    }
    vector<int> get(int x)
    {
        vector<int> ans;
        for (int i = 0; i < x; i++)
            ans.push_back(a[i]);
        return ans;
    }
    Node operator+(const Node &rhs) const
    {
        vector<int> ans;
        auto p = a.begin(), q = rhs.a.begin();
        while (ans.size() < 10 && p != a.end() && q != rhs.a.end())
        {
            if (*p < *q)
                ans.push_back(*p), ++p;
            else
                ans.push_back(*q), ++q;
        }
        while (ans.size() < 10 && p != a.end())
            ans.push_back(*p), ++p;
        while (ans.size() < 10 && q != rhs.a.end())
            ans.push_back(*q), ++q;
        return {ans};
    }
    Node &operator+=(const Node &rhs)
    {
        (*this) = (*this) + rhs;
    }
};

int fa[N][20], dep[N];
Node node[N][20];

void dfs(int p, int from)
{
    for (int q : g[p])
        if (q != from)
        {
            fa[q][0] = p;
            for (int i = 1; i < 20; i++)
                fa[q][i] = fa[fa[q][i - 1]][i - 1],
                node[q][i] = node[q][i - 1] + node[fa[q][i - 1]][i - 1];
            dep[q] = dep[p] + 1;
            dfs(q, p);
        }
}

Node query(int p, int q)
{
    Node ans;
    if (dep[p] < dep[q])
        swap(p, q);
    for (int i = 19; i >= 0; i--)
        if (dep[fa[p][i]] >= dep[q])
            ans += node[p][i], p = fa[p][i];
    for (int i = 19; i >= 0; i--)
        if (fa[p][i] != fa[q][i])
            ans += node[p][i], ans += node[q][i], p = fa[p][i], q = fa[q][i];
    if (p != q)
        ans += node[p][0], ans += node[q][0], p = fa[p][0], q = fa[q][0];
    ans += node[p][0];
    return ans;
}

signed main()
{
    ios::sync_with_stdio(false);
    cin >> n >> m >> q;
    for (int i = 1; i < n; i++)
    {
        int u, v, w;
        cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    for (int i = 1; i <= m; i++)
    {
        int x;
        cin >> x;
        node[x][0].insert(i);
    }
    dep[1] = 1;
    dfs(1, 0);

    for (int i = 1; i <= q; i++)
    {
        int u, v, a;
        cin >> u >> v >> a;
        auto ans = query(u, v);
        int sz = min(1ull * ans.a.size(), 1ull * a);
        cout << sz << " ";
        for (int i = 0; i < sz; i++)
            cout << ans.a[i] << " ";
        cout << endl;
    }
}