brz的树【牛客练习赛72 F】【虚树+dfs序可持久化线段树+树上差分】

题目链接

  有N个点的树，现在我们有M次询问，询问（u、v），只在u、v子树内出现的颜色的个数（u子树 并 v子树）。

  首先，可以将问题拆开来讨论，如果对于一个点的时候是怎样的呢？

  先求一个点时候的答案：

  那么，实际上就是求“颜色总数-不在子树内出现的颜色总数”。那么，实际上，我们可以将点数扩展成2N这么多个，然后查询出现在dfn[u] + siz[u] ~ dfn[u] + N - 1内的颜色的个数。当然，不嫌麻烦的话，就不需要这么做，直接写树剖，然后找到每个颜色的总的父亲节点，给1～point这条链全部“+1”，然后之后单点查询即可。——于是，我们就可以求得每个点的自身子树所产生的答案了。

  然后考虑两个点的时候，当且仅当我们对一个颜色构建虚树，虚树有虚根（虚根颜色不为该颜色），且虚根有且仅有两个虚子节点的时候，才需要考虑这种颜色的贡献。

  于是，求两个点时候的答案：

  我们不妨对虚子节点进行配对，譬如说“u指向v”这样的，于是，对于查询的x、y，我们实际上就是想知道x的子树内，有多少是指向y子树内的。那么，对于x的子树区间为dfn[x] ~ dfn[x] + siz[x] - 1，同理，y的子树。那么，比较显然的，这不就是用dfs序来构造一个主席树，然后在主席树上进行差分，找到x子树对于y子树的贡献吗？

  所以，我们可以利用虚树来找到这样的u、v对，然后我们让dfs序小的指向dfs序大的，最后，我们对它的dfs序上建立可持久化线段树，然后对于查询的x、y对，我们用dfs序小的作为根的差分，去查dfs序大的部分。——于是，利用这些，我们就可以知道这样的关系对对每个答案产生的贡献了。

#include <iostream>
#include <cstdio>
#include <cmath>
#include <string>
#include <cstring>
#include <algorithm>
#include <limits>
#include <vector>
#include <stack>
#include <queue>
#include <set>
#include <map>
#include <bitset>
#include <unordered_map>
#include <unordered_set>
#define lowbit(x) ( x&(-x) )
#define pi 3.141592653589793
#define e 2.718281828459045
#define INF 0x3f3f3f3f
#define HalF (l + r)>>1
#define lsn rt<<1
#define rsn rt<<1|1
#define Lson lsn, l, mid
#define Rson rsn, mid+1, r
#define QL Lson, ql, qr
#define QR Rson, ql, qr
#define myself rt, l, r
#define pii pair<int, int>
#define MP(a, b) make_pair(a, b)
using namespace std;
typedef unsigned long long ull;
typedef unsigned int uit;
typedef long long ll;
const int maxN = 1e5 + 7;
int N, M, head[maxN], cnt, col[maxN], lsan[maxN], _UP;
vector<int> vt[maxN];
struct Eddge
{
    int nex, to;
    Eddge(int a=-1, int b=0):nex(a), to(b) {}
} edge[maxN << 1];
inline void addEddge(int u, int v)
{
    edge[cnt] = Eddge(head[u], v);
    head[u] = cnt ++;
}
inline void _add(int u, int v) { addEddge(u, v); addEddge(v, u); }
inline void init()
{
    cnt = 0;
    for(int i = 1; i <= N; i ++) head[i] = -1;
}
int deep[maxN], fa[maxN][20], LOG_2[maxN], siz[maxN], dfn[maxN], idx, rid[maxN << 1];
bool cmp(int a, int b) { return dfn[a] < dfn[b]; }
void dfs(int u, int father)
{
    siz[u] = 1;
    fa[u][0] = father;
    deep[u] = deep[father] + 1;
    dfn[u] = ++idx;
    rid[idx] = u;
    for(int i = 0; i < 18; i ++) fa[u][i + 1] = fa[fa[u][i]][i];
    for(int i = head[u], v; ~i; i = edge[i].nex)
    {
        v = edge[i].to;
        if(v == father) continue;
        dfs(v, u);
        siz[u] += siz[v];
    }
}
int lca(int u, int v)
{
    if(deep[u] < deep[v]) swap(u, v);
    int det = deep[u] - deep[v];
    for(int i = LOG_2[det]; i >= 0; i --)
    {
        if((det >> i) & 1) u = fa[u][i];
    }
    if(u == v) return u;
    for(int i = LOG_2[deep[v]]; i >= 0; i --)
    {
        if(fa[u][i] ^ fa[v][i])
        {
            u = fa[u][i];
            v = fa[v][i];
        }
    }
    return fa[u][0];
}
int stk[maxN], top;
vector<int> son;
void Insert(int u)
{
    if(top <= 1)
    {
        stk[++ top] = u;
        return;
    }
    int p = lca(u, stk[top]);
    while(top >= 2 && dfn[p] <= dfn[stk[top - 1]])
    {
        if(top == 2) son.push_back(stk[top]);
        top --;
    }
    if(stk[top] ^ p)
    {
        stk[top] = p;
    }
    stk[++ top] = u;
}
vector<int> nex[maxN];
void solve(int op)
{
    sort(vt[op].begin(), vt[op].end(), cmp);
    int root = vt[op][0];
    for(int u : vt[op]) root = lca(root, u);
    top = 0; son.clear();
    Insert(root);
    for(int u : vt[op])
    {
        if(u == root)
        {
            top = 0;
            son.clear();
            break;
        }
        Insert(u);
    }
    if(top >= 2) son.push_back(stk[2]);
    if(son.size() == 2)
    {
        int x = son[0], y = son[1];
        if(dfn[x] > dfn[y]) swap(x, y);
        nex[dfn[x]].push_back(dfn[y]);
    }
}
int t[maxN << 1];
void add(int x, int v) { while(x <= (N << 1)) { t[x] += v; x += lowbit(x); } }
int sum(int x) { int res = 0; while(x) { res += t[x]; x -= lowbit(x); } return res; }
vector<pii> ques[maxN << 1];
int las_col[maxN] = {0};
int ans[maxN];
namespace Segement
{
    const int maxP = maxN * 30;
    int tree[maxP], lc[maxP], rc[maxP];
    int root[maxN], tot;
    void build(int &rt, int old, int l, int r, int qx)
    {
        rt = ++ tot;
        lc[rt] = lc[old]; rc[rt] = rc[old]; tree[rt] = tree[old] + 1;
        if(l == r) return;
        int mid = HalF;
        if(qx <= mid) build(lc[rt], lc[old], l, mid, qx);
        else build(rc[rt], rc[old], mid + 1, r, qx);
    }
    int query(int rl, int rr, int l, int r, int ql, int qr)
    {
        if(ql <= l && qr >= r) return tree[rr] - tree[rl];
        int mid = HalF;
        if(qr <= mid) return query(lc[rl], lc[rr], l, mid, ql, qr);
        else if(ql > mid) return query(rc[rl], rc[rr], mid + 1, r, ql, qr);
        else return query(lc[rl], lc[rr], l, mid, ql, qr) + query(rc[rl], rc[rr], mid + 1, r, ql, qr);
    }
}
using namespace Segement;
struct Question
{
    int u, v, id;
    Question(int a=0, int b=0, int c=0):u(a), v(b), id(c) {}
};
vector<Question> qt;
int main()
{
    for(int i = 2; i < maxN; i ++) LOG_2[i] = LOG_2[i >> 1] + 1;
    scanf("%d%d", &N, &M);
    init();
    for(int i = 1; i <= N; i ++) { scanf("%d", &col[i]); lsan[i] = col[i]; }
    sort(lsan + 1, lsan + N + 1);
    _UP = (int)(unique(lsan + 1, lsan + N + 1) - lsan - 1);
    for(int i = 1; i <= N; i ++)
    {
        col[i] = (int)(lower_bound(lsan + 1, lsan + _UP + 1, col[i]) - lsan);
        vt[col[i]].push_back(i);
    }
    for(int i = 1, u, v; i < N; i ++)
    {
        scanf("%d%d", &u, &v);
        _add(u, v);
    }
    dfs(1, 0);
    for(int i = 1; i <= N; i ++) rid[N + i] = rid[i];
    for(int i = 1; i <= _UP; i ++) solve(i);
    for(int i = 1, x, y, p, l, r; i <= M; i ++)
    {
        scanf("%d%d", &x, &y);
        if(dfn[x] > dfn[y]) swap(x, y);
        p = lca(x, y);
        if(p == x)
        {
            ans[i] = _UP;
            l = dfn[p] + siz[p]; r = N + dfn[p] - 1;
            ques[r].push_back(MP(l, i));
        }
        else
        {
            ans[i] = _UP << 1;
            l = dfn[x] + siz[x]; r = N + dfn[x] - 1;
            ques[r].push_back(MP(l, i));
            l = dfn[y] + siz[y]; r = N + dfn[y] - 1;
            ques[r].push_back(MP(l, i));
            qt.push_back(Question(x, y, i));
        }
    }
    for(int i = 1, u, c, id, lx; i <= (N << 1); i ++)
    {
        u = rid[i];
        c = col[u];
        add(i, 1);
        if(las_col[c]) add(las_col[c], -1);
        las_col[c] = i;
        for(pii it : ques[i])
        {
            id = it.second;
            lx = it.first;
            ans[id] -= sum(i) - sum(lx - 1);
        }
    }
    for(int i = 1; i <= N; i ++)
    {
        root[i] = root[i - 1];
        for(int j : nex[i])
        {
            build(root[i], root[i], 1, N, j);
        }
    }
    for(Question it : qt)
    {
        int u = it.u, v = it.v, id = it.id;
        int l = dfn[u], r = dfn[u] + siz[u] - 1, ql = dfn[v], qr = dfn[v] + siz[v] - 1;
        ans[id] += query(root[l - 1], root[r], 1, N, ql, qr);
    }
    for(int i = 1; i <= M; i ++) printf("%d
", ans[i]);
    return 0;
}