UVA12424 Answering Queries on a Tree

题目大意

给定一棵 $N$ 个结点的树，每个结点有一个颜色 $C_i$。

有 $M$ 次操作，每次操作分为 $2$ 种：

0 u c：将 $C_u$ 修改为 $c$。

1 u v：输出 $u$ 到 $v$ 的路径中出现最多的颜色的出现次数。

对于 $100\%$ 的数据，满足 $1\leq N,M\leq 10^5,\forall i \in[1,N],1\leq C_i \leq 10$。

解题思路

不难想到树剖。

因为颜色最多只有 $10$ 种，那么可以每种颜色维护一下出现的位置（在某个出现记为 $1$，否则为 $0$，用线段树或树状数组都行。

那么 $u$ 到 $v$ 的路径中 $c$ 颜色的出现次数，即为 $c$ 颜色对应树状数组上的和。

对于操作 $0$，将结点在原来的颜色的那棵树中对应的位置减一，将结点在新的颜色的那棵树中对应的位置减一加一。

对于操作 $1$，暴力枚举 $10$ 种颜色，统计 $u$ 到 $v$ 的路径中每种颜色的出现次数再取最大值。

时间复杂度 $\mathcal O(n+m \log^2 n)$。

具体实现见代码。

CODE

#include <bits/stdc++.h>

using namespace std;

inline int read()
{
    int x = 0, f = 1;
    char c = getchar();
    while (c < '0' || c > '9')
    {
        if (c == '-')
            f = -1;
        c = getchar();
    }
    while (c >= '0' && c <= '9')
    {
        x = x * 10 + c - '0';
        c = getchar();
    }
    return x * f;
}

#define js(x) memset(x, 0, sizeof x)

const int _ = 3e5;

int T, n, m;

int col[_];

vector<int> d[_];

int cntnode, dfn[_], siz[_], top[_], fa[_], dep[_], hson[_];

void dfs1(int u, int D = 1)
{
    dep[u] = D;
    siz[u] = 1;
    for (int v : d[u])
    {
        if (siz[v])
            continue;
        fa[v] = u;
        dfs1(v, D + 1);
        siz[u] += siz[v];
        if (siz[v] > siz[hson[u]])
            hson[u] = v;
    }
}

void dfs2(int u, int topf)
{
    top[u] = topf;
    dfn[u] = ++cntnode;
    if (!hson[u])
        return;
    dfs2(hson[u], topf);
    for (int v : d[u])
    {
        if (top[v])
            continue;
        dfs2(v, v);
    }
}

struct BIT
{
    int c[_];
    int lowbit(int x)
    {
        return x & -x;
    }
    void update(int x, int d)
    {
        for (int i = x; i <= n; i += lowbit(i))
            c[i] += d;
    }
    int query(int x)
    {
        int res = 0;
        for (int i = x; i; i -= lowbit(i))
            res += c[i];
        return res;
    }
    int Query(int u, int v)
    {
        int res = 0;
        while (top[u] != top[v])
        {
            if (dep[top[u]] < dep[top[v]])
                swap(u, v);
            res += query(dfn[u]) - query(dfn[top[u]] - 1);
            u = fa[top[u]];
        }
        if (dep[u] > dep[v])
            swap(u, v);
        res += query(dfn[v]) - query(dfn[u] - 1);
        return res;
    }
} bit[11];

void init()
{
    cntnode = 0;
    js(dfn);
    js(siz);
    js(top);
    js(fa);
    js(hson);
    js(dep);
    for (int i = 1; i <= 10; ++i)
        js(bit[i].c);
    for (int i = 0; i < _; ++i)
        d[i].clear();
}

signed main()
{
    T = read();
    while (T--)
    {
        init();
        n = read(), m = read();
        for (int i = 1; i <= n; ++i)
            col[i] = read();
        for (int i = 1, u, v; i < n; ++i)
        {
            u = read(), v = read();
            d[u].push_back(v);
            d[v].push_back(u);
        }
        dfs1(1);
        dfs2(1, 1);
        for (int i = 1; i <= n; ++i)
            bit[col[i]].update(dfn[i], 1);
        while (m--)
        {
            int opt = read(), u = read(), v = read();
            if (opt == 0)
            {
                bit[col[u]].update(dfn[u], -1);
                col[u] = v;
                bit[col[u]].update(dfn[u], 1);
            }
            else
            {
                int ans = 0;
                for (int i = 1; i <= 10; ++i)
                    ans = max(ans, bit[i].Query(u, v));
                printf("%d\n", ans);
            }
        }
    }
    return 0;
}