牛客练习赛72 F brz的树

这场题好妙啊，不仅打开我的凸包新大门，这两题树也真是妙阿 brz， yyds！

题意:

在一颗点被染色的带根树上询问仅出现在两颗子树中的颜色数

题解:

颜色数有两种贡献，一种仅在单颗子树中出现的颜色，和仅在两颗子树中出现的颜色

两颗子树也有两种情况，一种是包含关系，对于这种直接计算第一种贡献即可

第二种两颗子树的情况，先计算第一颗和第二颗的第一种贡献相加

考虑第二种贡献，设a和b为保证颜色C仅在其子树中存在的dfs序最大的两个节点

容易想到，要么存在这两个节点，要么颜色C对第二种没有贡献

然后是极妙的地方了， 对颜色C求虚树，直接看根节点有没有两个儿子即可。同时对所有儿子求lca，可得到颜色C仅在其子树中出现的dfs序最大的点，记录一下即可

最后计算贡献时，对于ab链，经过a点时，将b点的值加一，对询问点的子树求和即可计算出第二种贡献

完整代码

struct BitTree {
    int res[Ma];
    void add(int pos, int val) {for (; pos < Ma; pos += lowbit(pos)) res[pos] += val;}
    int sum(int pos) {int s = 0; for (; pos; pos -= lowbit(pos)) s += res[pos]; return s;}
    int range(int l, int r) {return sum(r - 1) - sum(l - 1);}
};

int c[Ma];
vector<int> val[Ma];
int ans[Ma];
typedef pair<int, int*> ask_t;
vector<ask_t> ak[Ma];

template<typename T>
struct HLD {
    BitTree bt;
    static const int Ma = 1e5 + 100;
    int fa[Ma], dep[Ma], siz[Ma], son[Ma], st[Ma];
    int top[Ma], dfn[Ma], rnk[Ma]; int cnt;
    typedef std::vector<std::vector<T>> Tree;
    const std::vector<std::vector<T> >& g;

    void dfs1(int u) {
        son[u] = -1; siz[u] = 1;
        for (const auto& i : g[u]) if (!dep[i]) {
            dep[i] = dep[u] + 1;
            fa[i] = u;
            dfs1(i);
            siz[u] += siz[i];
            if (son[u] == -1 or siz[i] > siz[son[u]]) son[u] = i;
        }
    }
    void dfs2(int u, int f) {
        top[u] = f; dfn[u] = ++cnt; rnk[cnt] = u;
        if (!~son[u]) return ;
        dfs2(son[u], f);
        for (const auto& i : g[u]) if (i != son[u] and i != fa[u])
            dfs2(i, i);
    }
    HLD(std::vector<std::vector<T> >& gg) : g(gg) {
        memset(dep, 0, sizeof dep);
        dep[0] = 1; cnt = 0;
        dfs1(0); dfs2(0, 0);
    }
    int LCA(int u, int v) const {
        while (top[u] != top[v]) {
            if (dep[top[u]] > dep[top[v]]) u = fa[top[u]];
            else v = fa[top[v]];
        }
        return dep[u] > dep[v] ? v : u;
    }

    void solve(int u, int fa) {
        rap (i, ak[u]) *i.S -= bt.range(dfn[i.F], dfn[i.F] + siz[i.F]);
        rap (i, val[u]) bt.add(dfn[i], 1);
        if (~son[u]) solve(son[u], u);
        rap (i, g[u]) if (i != fa and i != son[u]) solve(i, u);
        rap (i, ak[u]) *i.S += bt.range(dfn[i.F], dfn[i.F] + siz[i.F]);
    }

    Tree& vtree(std::vector<int>& key) {
        static Tree tr(cnt); int top = 0; st[top] = 0; tr[0].clear();
        cc_hash_table<int, null_type> vis;
        std::sort(key.begin(), key.end(), [&](int a, int b) {return dfn[a] < dfn[b];});
        static std::function<void(int, int)> add = [&] (const int& u, const int& v) {
            if (vis.find(u) == vis.end()) tr[u].clear(), vis.insert(u);
            tr[u].emplace_back(v);
        };
        static std::function<void(int)> insert = [&] (int x) {
            if (top == 0) {x != 0 ? st[++top] = x : 0; return ;}
            int lca = LCA(x, st[top]);
            if (lca == st[top]) return ; // st[++top] = x;
            else {
                while (top > 0 and dfn[st[top - 1]] >= dfn[lca])
                    add(st[top - 1], st[top]), --top;
                if (lca != st[top]) add(lca, st[top]), st[top] = lca;
                st[++top] = x;
            }
        };
        for (const auto& i : key) insert(i);
        while (top > 0) add(st[top - 1], st[top]), --top;
        return tr;
    }
};

int cc[Ma];
int own[Ma];
vector< vector<int> > g;
vector<int> vp[Ma];

void cal(int u, int fa) {
    rap (i, g[u]) if (i != fa) cal(i, u);
    if (u) own[fa] += own[u];
}

signed main() {
    #if SYNC==0
    ios::sync_with_stdio(false);
    cin.tie(0);
    #endif
    int n, m; cin >> n >> m;
    rep (i, n) {
        cin >> c[i];
        vp[c[i]].ep(i);
    }
    memcpy(cc, c, sizeof c);
    sort(cc, cc + n);
    int len = unique(cc, cc + n) - cc;
    g.resize(n);
    rep (i, n - 1) {
        int u, v; cin >> u >> v; --u, --v;
        g[u].ep(v); g[v].ep(u);
    }
    static HLD<int> hld(g);
    rep (i, len) {
        vector< vector<int> >& tp = hld.vtree(vp[cc[i]]);
        int aim = 0;
        if (tp[0].size() == 1 and c[tp[0][0]] != cc[i]) aim = tp[0][0];
        debug(cc[i], aim, tp[aim].size());
        if (cc[i] != c[0] and tp[aim].size() == 2) {
            int a = tp[aim][0], b = tp[aim][1];
            debug(a, b);
            if (hld.dfn[a] > hld.dfn[b]) swap(a, b);
            val[a].ep(b);
        }
        if (cc[i] == c[0] or tp.empty()) ++own[0];
        else {
            int now = tp[0][0];
            for (size_t i = 1; i < tp[0].size(); i++) now = hld.LCA(now, tp[0][i]);
            ++own[now];
        }
    }
    cal(0, 0);
    debug(own[1], own[2]);
    rep (i, m) {
        int x, y; cin >> x >> y; --x, --y;
        if (hld.dfn[x] > hld.dfn[y]) swap(x, y);
        int lca = hld.LCA(x, y);
        if (lca == x) {ans[i] = own[x]; continue;}
        else ans[i] = own[x] + own[y];
        ak[x].ep(mkp(y, ans + i));
    }
    hld.solve(0, 0);
    rep (i, m) cout << ans[i] << endl;

    return 0;
}