天下一 2015 本戦 F - 根付き木のみさわさん

オイラーツアーの練習に解いたけど、発想も面白いん！！！

問題へのリンク

問題概要

頂点 1 を根とする、頂点数 N の根付き木がある。以下の Q 個のクエリに答えよ:

• M 個の頂点 v1, v2, ..., vM に実をつける。「その頂点を根とする部分木に含まれる実の個数が K 個以上」という条件を満たす頂点のうち、根からの深さの最大値を出力せよ。

制約

• 3 <= N <= 10⁵
• 1 <= Q <= 10⁵
• Q 個の M の総和は 10⁵ 以下

考えたこと

「部分木に含まれる実の個数」というのがいかにもオイラーツアーっぽい。そうすると、オイラーツアーに BIT を乗せることで

• 点 v を根とした部分木に含まれる実の個数

は O(log N) でとってこれるようになる。問題はここからで、愚直には全頂点についてその部分木に含まれる木の実の個数を数えて、それが K 以上になっているものについての深さの最大値を求めればよい。このままでは各クエリに O(N logN) かかってしまうので高速化したい (これ全体の作業自体はツリー DP で O(N) でもできる)。

ここで、「各深さについての、その深さに含まれる全頂点についての、その部分木の木の実の個数の最大値」みたいなのを管理したくなると、手が出しづらくなる。HL 分解などの方針も考えられるが、上手い解法があった。

まず、木の実をついた頂点どうしの LCA のみを考えればよいことがわかる。しかしこのままではまだ二乗のオーダーがかかる。この時点で僕はこの方向性の考察を諦めてしまった。さらに突き進むと、

• オイラーツアー上の木の実の頂点のうち連続する K 個についての、その両端の LCA のみを考えればよい

ということがわかる。ここまでくると、なんと全クエリ総合しても考えるべき頂点数は O(M) になっている！！！！！

というわけで、EulerTour + LCA + BIT で O(N + MlogN) でできる。

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

template<class VAL> struct RMQ {
    vector<pair<VAL, int> > dat;
    int SIZE_R;
    VAL INF = 1<<29; // to be set
    
    RMQ(int n = 110000) { init(n); }
    void init(int n) {
        SIZE_R = 1;
        while (SIZE_R < n) SIZE_R *= 2;
        dat.resize(SIZE_R * 2 - 1);
        for (int i = 0; i < (int)dat.size(); ++i) dat[i] = make_pair(INF, -1);
    }

    inline void set(int a, VAL v) {
        int k = a + SIZE_R - 1;
        dat[k] = make_pair(v, a);
        while (k > 0) {
            k = (k-1)/2;
            dat[k] = min(dat[k*2+1], dat[k*2+2]);
        }
    }
    
    inline pair<VAL,int> get(int a, int b, int k, int l, int r) {
        if (r <= a || b <= l) return make_pair(INF, -1);
        if (a <= l && r <= b) return dat[k];
        else {
            pair<VAL,int> vl = get(a, b, k*2+1, l, (l+r)/2);
            pair<VAL,int> vr = get(a, b, k*2+2, (l+r)/2, r);
            return min(vl, vr);
        }
    }
    inline pair<VAL,int> get(int a, int b) { return get(a, b, 0, 0, SIZE_R); }
    
    void print() {
        for (int i = 0; i < SIZE_R; ++i) {
            VAL val = (*this)[i];
            if (val < INF) cout << val;
            else cout << "INF";
            if (i != SIZE_R-1) cout << ",";
        }
        cout << endl;
    }
};

template<class Abel> struct BIT {
    vector<Abel> dat;
    Abel UNITY_SUM = 0;                        // to be set
    
    /* [1, n] */
    BIT(int n = 110000) { init(n); }
    void init(int n) {
        dat.resize(n + 1);
        for (int i = 0; i < (int)dat.size(); ++i) dat[i] = UNITY_SUM;
    }
    
    /* a is 1-indexed */
    inline void add(int a, Abel x) {
        for (int i = a; i < (int)dat.size(); i += i & -i)
            dat[i] = dat[i] + x;
    }
    
    /* [1, a], a is 1-indexed */
    inline Abel sum(int a) {
        Abel res = UNITY_SUM;
        for (int i = a; i > 0; i -= i & -i)
            res = res + dat[i];
        return res;
    }
    
    /* [a, b), a and b are 1-indexed */
    inline Abel sum(int a, int b) {
        return sum(b - 1) - sum(a - 1);
    }
    
    void print() {
        for (int i = 1; i < (int)dat.size(); ++i) cout << sum(i, i + 1) << ",";
        cout << endl;
    }
};

typedef vector<vector<int> > Graph;

struct EulerTour {
    // main results
    vector<int> depth;
    vector<int> node; // the node-number of i-th element of Euler-tour
    vector<int> vf, ve; // the index of Euler-tour of node v
    
    // sub results
    RMQ<int> rmq; // depth (to find LCA)
    BIT<int> bit; // to calc sum of sub-tree
    
    EulerTour(const Graph &tree) { init(tree); }
    void init(const Graph &tree) {
        int V = (int)tree.size();
        depth.resize(V*2-1); node.resize(V*2-1); vf.resize(V); ve.resize(V);
        rmq.init(V*2-1); bit.init(V*2-1);
        int k = 0;
        dfs(tree, 0, -1, 0, k);
        for (int i = 0; i < V*2-1; ++i) rmq.set(i, depth[i]);
    }
    
    void dfs(const Graph &tree, int v, int par, int dep, int &ord) {
        node[ord] = v;
        depth[ord] = dep;
        vf[v] = ve[v] = ord;
        ++ord;
        for (int i = 0; i < tree[v].size(); ++i) {
            auto e = tree[v][i];
            if (e != par) {
                dfs(tree, e, v, dep+1, ord);
                node[ord] = v;
                depth[ord] = dep;
                ve[v] = ord;
                ++ord;
            }
        }
    }

    int LCA(int u, int v) {
        int a = vf[u], b = vf[v];
        if (a > b) swap(a, b);
        return node[rmq.get(a, b+1).second];
    }
};

int N, Q;
Graph tree;

int main() {
    cin >> N;
    tree.clear(); tree.resize(N);
    for (int i = 0; i < N-1; ++i) {
        int a, b; cin >> a >> b; --a, --b;
        tree[a].push_back(b); tree[b].push_back(a);
    }
    EulerTour et(tree);
    cin >> Q;
    for (int q = 0; q < Q; ++q) {
        int M, K; cin >> M >> K;
        vector<int> ids(M);
        for (int i = 0; i < M; ++i) {
            int v; cin >> v; --v;
            int vf = et.vf[v];
            ids[i] = vf;
            et.bit.add(vf+1, 1);
        }
        sort(ids.begin(), ids.end());
        int res = 0;
        for (int i = 0; i+K-1 < M; ++i) {
            int iu = ids[i], iv = ids[i+K-1];
            int u = et.node[iu], v = et.node[iv];
            int lca = et.LCA(u, v);
            int depth = et.depth[et.vf[lca]];
            res = max(res, depth);
        }
        cout << res << endl;
    }
}