Codeforces Round #326 (Div. 1) C. Duff in the Army

問題

codeforces.com

解法

lca を使って lca と同じようなテクニックを使って解きます。

lca の時と同じように, vals[logk][v] = (v から 2^logk だけ上までの頂点における, ID の小さい順の数列 p1, p2, ..., p10) というのを保持しておきます。
で, 頂点 u と 頂点 v の間の path にある ID を求めるときは, u と v の lca を求めて, u から lca, v から lca の間の ID 数列を merge させて小さい順に並べたものを答えとすれば良いです。

class Tree {
public:
    Tree(int V, int root) : V(V), root(root) {
        T.resize(V);
        for (int i = 0; i < MAXLOGV; i++) {
            parent[i].resize(V);
            vals[i].resize(V);
        }
        depth.resize(V);
    }
    // uとvをつなぐ
    // lcaを求めることが主目的なので無向グラフとしている
    void unite(int u, int v) {
        T[u].push_back(v);
        T[v].push_back(u);
    }
    // initする
    // コンストラクタだけじゃなくてこれも呼ばないとlcaが求められないぞ
    void init() {
        dfs(root, 0, 0);
    }
    // uとvのlcaを求める
    int lca(int u, int v) const {
        if (depth[u] > depth[v]) swap(u, v);
        for (int k = 0; k < MAXLOGV; k++) {
            if ((depth[v] - depth[u])>>k & 1) {
                v = parent[k][v];
            }
        }
        if (u == v) return u;
        for (int k = MAXLOGV-1; k >= 0; k--) {
            if (parent[k][u] != parent[k][v]) {
                u = parent[k][u];
                v = parent[k][v];
            }
        }
        return parent[0][u];
    }
    // uとvの距離を求める
    // edgeを定義しないといけない時はこれじゃダメ
    int dist(int u, int v) const {
        int p = lca(u, v);
        return (depth[u]-depth[p]) + (depth[v]-depth[p]);
    }
    void dfs(int v, int p, int d) {
        parent[0][v] = p;
        depth[v] = d;
        for (int i = 1; i < MAXLOGV; i++) {
            parent[i][v] = parent[i-1][parent[i-1][v]];
            vals[i][v] = merge(vals[i-1][v], vals[i-1][parent[i-1][v]]);
        }
        for (int next : T[v]) {
            if (next != p) dfs(next, v, d+1);
        }
    }
    static const int MAXLOGV = 25;
    // グラフの隣接リスト表現
    vector<vector<int> > T;
    // 頂点の数
    int V;
    // 根ノードの番号
    int root;

    // 親ノード
    vector<int> parent[MAXLOGV];
    // 根からの深さ
    vector<int> depth;

    struct node {
        int a[11];
        node() {
            memset(a, 63, sizeof(a));
        }
        void insert(int val) {
            a[10] = val;
            sort(a, a+11);
        }
    };
    vector<node> vals[MAXLOGV];
    node merge(node x, node y) {
        node ans = x;
        for (int i = 0; i < 11; i++) {
            ans.insert(y.a[i]);
        }
        return ans;
    }
    node get(int v, int k) {
        node ans;
        for (int i = 0; i < MAXLOGV; i++) {
            if ((1<<i)&k) {
                ans = merge(ans, vals[i][v]);
                v = parent[i][v];
            }
        }
        return ans;
    }
};

int main() {
    int n, m, q;
    scanf("%d%d%d", &n, &m, &q);
    Tree tree(n, 0);
    for (int i = 0; i < n-1; i++) {
        int a, b;
        scanf("%d%d", &a, &b);
        tree.unite(a-1, b-1);
    }
    for (int i = 0; i < m; i++) {
        int c;
        scanf("%d", &c);
        tree.vals[0][c-1].insert(i+1);
    }
    tree.init();
    while (q--) {
        int v, u, k;
        scanf("%d%d%d", &v, &u, &k);
        v--; u--;
        int lca = tree.lca(u, v);
        Tree::node x = tree.get(v, tree.depth[v]-tree.depth[lca]);
        Tree::node y = tree.get(u, tree.depth[u]-tree.depth[lca]+1);
        Tree::node ans = tree.merge(x, y);
        int tmp = 0;
        while (tmp < k && ans.a[tmp] <= m) tmp++;
        k = tmp;
        printf("%d ", k);
        for (int i = 0; i < k; i++) {
            printf("%d ", ans.a[i]);
        }
        printf("\n");
    }
    return 0;
}