题解：树上计数2

一、题目分析

本题给定一棵有(N)个节点的树，每个节点都有一个整数权值，需要处理(M)次询问，每次询问两个节点(u)和(v)，要求输出从(u)到(v)路径上（包括两端点）不同点权值的种类数。

二、解题思路

本题主要通过树的欧拉序、最近公共祖先（LCA）算法结合莫队算法来解决。先将树转化为欧拉序，把树上路径问题转化为序列区间问题，再利用莫队算法处理区间询问。

（一）数据结构与变量定义

const int N = 100010;
int n, m, len;
int w[N];
int h[N], e[N], ne[N], idx;
int depth[N], f[N][16];
int seq[N], top, first[N], last[N];
int cnt[N], st[N], ans[N];
int que[N];
struct Query
{
    int id, l, r, p;
}q[N];
vector<int> nums;

• N：定义节点数量和询问数量的最大值。
• n：存储树的节点数。
• m：存储询问的数量。
• len：存储分块的长度。
• w[N]：数组，用于存储每个节点的权值。
• h[N]、e[N]、ne[N]、idx：用于邻接表存储树的边信息，h是表头数组，e存储边的终点，ne指向下一条边的索引，idx是边的计数。
• depth[N]：数组，用于存储每个节点的深度。
• f[N][16]：二维数组，用于倍增法求最近公共祖先，f[u][k]表示节点u向上跳(2^k)步后的祖先节点。
• seq[N]：数组，存储树的欧拉序。
• top：表示欧拉序数组seq的当前索引。
• first[N]和last[N]：数组，分别记录节点在欧拉序中第一次和最后一次出现的位置。
• cnt[N]：数组，用于记录每种权值在当前区间内出现的次数。
• st[N]：数组，标记节点是否在当前处理的区间内（通过异或操作实现添加和删除效果）。
• ans[N]：数组，用于存储每个询问的答案。
• que[N]：数组，用于广度优先搜索（BFS）时的队列。
• Query结构体：用于存储每个询问的信息，包括询问的编号id，区间左端点l，右端点r，以及路径的最近公共祖先节点p。
• nums：向量，用于存储需要离散化处理的节点权值。

（二）辅助函数

1. 添加边函数add_edge：

void add_edge(int a, int b)
{
    e[idx] = b, ne[idx] = h[a], h[a] = idx ++ ;
}

向邻接表中添加一条从节点a到节点b的边。

1. 深度优先搜索函数dfs：

void dfs(int u, int father)
{
    seq[ ++ top] = u;
    first[u] = top;
    for (int i = h[u]; ~i; i = ne[i])
    {
        int j = e[i];
        if (j != father) dfs(j, u);
    }
    seq[ ++ top] = u;
    last[u] = top;
}

对树进行深度优先搜索，构建树的欧拉序，并记录每个节点在欧拉序中第一次和最后一次出现的位置。

1. 广度优先搜索函数bfs：

void bfs()
{
    memset(depth, 0x3f, sizeof depth);
    depth[0] = 0, depth[1] = 1;
    int hh = 0, tt = 0;
    que[0] = 1;
    while (hh <= tt)
    {
        int t = que[hh ++ ];
        for (int i = h[t]; ~i; i = ne[i])
        {
            int j = e[i];
            if (depth[j] > depth[t] + 1)
            {
                depth[j] = depth[t] + 1;
                f[j][0] = t;
                for (int k = 1; k <= 15; k ++ )
                    f[j][k] = f[f[j][k - 1]][k - 1];
                que[ ++ tt] = j;
            }
        }
    }
}

使用广度优先搜索计算每个节点的深度，并预处理出倍增法求最近公共祖先所需的信息。

1. 最近公共祖先函数lca：

int lca(int a, int b)
{
    if (depth[a] < depth[b]) swap(a, b);
    for (int k = 15; k >= 0; k -- )
        if (depth[f[a][k]] >= depth[b])
            a = f[a][k];
    if (a == b) return a;
    for (int k = 15; k >= 0; k -- )
        if (f[a][k] != f[b][k])
        {
            a = f[a][k];
            b = f[b][k];
        }
    return f[a][0];
}

利用倍增法求出两个节点a和b的最近公共祖先。

1. 获取元素所属块的编号函数get：

int get(int x)
{
    return x / len;
}

根据元素的下标x，返回其所属块的编号。

1. 询问排序比较函数cmp：

bool cmp(const Query& a, const Query& b)
{
    int i = get(a.l), j = get(b.l);
    if (i != j) return i < j;
    return a.r < b.r;
}

用于对询问进行排序。先按区间左端点所属块的编号排序，若块编号相同，则按区间右端点排序。

1. 区间元素处理函数add：

void add(int x, int& res)
{
    st[x] ^= 1;
    if (st[x] == 0)
    {
        cnt[w[x]] -- ;
        if (!cnt[w[x]]) res -- ;
    }
    else
    {
        if (!cnt[w[x]]) res ++ ;
        cnt[w[x]] ++ ;
    }
}

在当前区间添加或删除一个节点x。通过异或操作标记节点是否在区间内，更新该节点权值的出现次数，并相应地更新不同权值的种类数res。

（三）主函数main

int main()
{
    scanf("%d%d", &n, &m);
    for (int i = 1; i <= n; i ++ ) scanf("%d", &w[i]), nums.push_back(w[i]);
    sort(nums.begin(), nums.end());
    nums.erase(unique(nums.begin(), nums.end()), nums.end());
    for (int i = 1; i <= n; i ++ )
        w[i] = lower_bound(nums.begin(), nums.end(), w[i]) - nums.begin();

    memset(h, -1, sizeof h);
    for (int i = 0; i < n - 1; i ++ )
    {
        int a, b;
        scanf("%d%d", &a, &b);
        add_edge(a, b), add_edge(b, a);
    }

    dfs(1, -1);
    bfs();

    for (int i = 0; i < m; i ++ )
    {
        int a, b;
        scanf("%d%d", &a, &b);
        if (first[a] > first[b]) swap(a, b);
        int p = lca(a, b);
        if (a == p) q[i] = {i, first[a], first[b]};
        else q[i] = {i, last[a], first[b], p};
    }

    len = sqrt(top);
    sort(q, q + m, cmp);

    for (int i = 0, L = 1, R = 0, res = 0; i < m; i ++ )
    {
        int id = q[i].id, l = q[i].l, r = q[i].r, p = q[i].p;
        while (R < r) add(seq[ ++ R], res);
        while (R > r) add(seq[R -- ], res);
        while (L < l) add(seq[L ++ ], res);
        while (L > l) add(seq[ -- L], res);
        if (p) add(p, res);
        ans[id] = res;
        if (p) add(p, res);
    }

    for (int i = 0; i < m; i ++ ) printf("%d\n", ans[i]);

    return 0;
}

1. 读取树的节点数n和询问数量m，读取每个节点的权值，将其存入nums向量用于离散化处理。对nums进行排序和去重，然后将w数组中的节点权值离散化。
2. 初始化邻接表，读取树的边信息，构建树的邻接表表示。
3. 对树进行深度优先搜索构建欧拉序，再进行广度优先搜索预处理最近公共祖先信息。
4. 读取每个询问的两个节点，求出它们的最近公共祖先，根据情况将询问转化为欧拉序上的区间，并存储到q数组中。
5. 计算分块的长度len，对询问按照cmp函数定义的规则进行排序。
6. 初始化双指针L和R以及答案变量res，按序处理每个询问：
   • 通过移动双指针L和R，调用add函数更新区间内不同权值的种类数res。
   • 如果路径的最近公共祖先p存在，对其进行相应处理（添加和恢复，以保证计算正确）。
   • 将当前询问的答案存储到ans数组中。
7. 输出每个询问的答案。

四、时间复杂度和空间复杂度

（一）时间复杂度

• 初始化操作：读取数据、离散化处理、构建邻接表，时间复杂度为(O(n + n\log n))。
• 深度优先搜索和广度优先搜索操作：时间复杂度为(O(n))。
• 处理询问操作：分块处理询问，使用双指针移动计算答案，时间复杂度为(O(m\sqrt{n}))。
• 总的时间复杂度为(O(n + n\log n + m\sqrt{n}))，近似为(O(m\sqrt{n}))（当n和m规模相近时）。

（二）空间复杂度

需要存储节点权值w[N]、邻接表信息、深度数组depth[N]、倍增数组f[N][16]、欧拉序数组seq[N]、询问信息q[N]、权值计数数组cnt[N]、节点标记数组st[N]、答案数组ans[N]以及离散化用的nums向量等，空间复杂度为(O(N + N + N + N\times16 + N + N + N + N + N)=O(N)) 。