[CF343D] Water Tree

Description

给定一棵树，提供 (3) 种操作，每次指定一个点，将它的所有孩子染色为 (1)，或它到根的路径染色为 (0)，或询问它是否被染色。

Solution

复习一下树链剖分。主要是关于询问的部分。

我们按照这样的方式去处理一个询问：如果 top[p] 和 top[q] 不相等，那么我们就从 top 深度较大的那个（设为 (p)）向上染一段 ([top[p],p])，然后将 (p) 修改为 (fa[top[p]])。最后对于剩下的，深度较大的 (x) 和深度较小的 (y)，我们还需要染一段 ([y,x])。所有操作依据 DFS 序进行。

线段树需要支持区间覆盖和单点询问，当然也可以转化为差分的单点修改区间求和。

#include <bits/stdc++.h>
using namespace std;

const int N = 2e6 + 5;

int n, m, t1, t2, t3;

int a[N], tag[N];

void pushup(int p)
{
    a[p] = a[p * 2] + a[p * 2 + 1];
}

void pushdown(int p, int l, int r)
{
    if (tag[p])
    {
        if (tag[p] == 1)
        {
            a[p * 2] = 1;
            a[p * 2 + 1] = 1;
            tag[p * 2] = 1;
            tag[p * 2 + 1] = 1;
        }
        else
        {
            a[p * 2] = 0;
            a[p * 2 + 1] = 0;
            tag[p * 2] = -1;
            tag[p * 2 + 1] = -1;
        }
        tag[p] = 0;
    }
}

void modify(int p, int l, int r, int ql, int qr, int val)
{
    if (l > qr || r < ql)
        return;
    if (l >= ql && r <= qr)
    {
        if (val)
        {
            a[p] = r - l + 1;
            tag[p] = 1;
        }
        else
        {
            a[p] = 0;
            tag[p] = -1;
        }
    }
    else
    {
        pushdown(p, l, r);
        modify(p * 2, l, (l + r) / 2, ql, qr, val);
        modify(p * 2 + 1, (l + r) / 2 + 1, r, ql, qr, val);
        pushup(p);
    }
}

int query(int p, int l, int r, int pos)
{
    int mid = (l + r) / 2;
    if (l == r)
    {
        return a[p];
    }
    else
    {
        pushdown(p, l, r);
        if (pos <= mid)
        {
            return query(p * 2, l, (l + r) / 2, pos);
        }
        else
        {
            return query(p * 2 + 1, (l + r) / 2 + 1, r, pos);
        }
    }
}

// void modify(int x,int y,int z,int ql,int qr,int val)
// {
//     for(int i=ql;i<=qr;i++) a[i]=val;
// }

// int query(int x,int y,int z,int ql)
// {
//     return a[ql];
// }

int ind, dfn[N], fin[N], top[N], fa[N], wson[N], dep[N], siz[N];

vector<int> g[N];

void presolve()
{
    function<void(int, int)> dfs1 = [&](int p, int f) -> void {
        siz[p] = 1;
        siz[0] = 0;
        for (int q : g[p])
        {
            if (q == f)
                continue;
            fa[q] = p;
            dep[q] = dep[p] + 1;
            dfs1(q, p);
            siz[p] += siz[q];
            if (siz[q] > siz[wson[p]])
            {
                wson[p] = q;
            }
        }
    };
    function<void(int, int)> dfs2 = [&](int p, int f) -> void {
        dfn[p] = ++ind;
        if (wson[p])
        {
            int q = wson[p];
            top[q] = top[p];
            dfs2(q, p);
        }
        for (int q : g[p])
        {
            if (q == f || q == wson[p])
                continue;
            top[q] = q;
            dfs2(q, p);
        }
        fin[p] = ind;
    };
    dep[1] = 1;
    dfs1(1, 0);
    top[1] = 1;
    dfs2(1, 0);
}

int tree_query(int p)
{
    return query(1, 1, n, dfn[p]);
}

void link_modify(int p, int q, int v)
{
    while (top[p] != top[q])
    {
        if (dep[top[p]] < dep[top[q]])
            swap(p, q);
        modify(1, 1, n, dfn[top[p]], dfn[p], v);
        p = fa[top[p]];
    }
    if (dep[p] < dep[q])
        swap(p, q);
    modify(1, 1, n, dfn[q], dfn[p], v);
}

void subtree_modify(int p, int v)
{
    modify(1, 1, n, dfn[p], dfn[p] + siz[p] - 1, v);
}

int main()
{
    ios::sync_with_stdio(false);

    cin >> n;
    for (int i = 1; i < n; i++)
    {
        cin >> t1 >> t2;
        g[t1].push_back(t2);
        g[t2].push_back(t1);
    }

    presolve();

    cin >> m;

    for (int i = 1; i <= m; i++)
    {
        cin >> t1 >> t2;
        if (t1 == 1)
        {
            subtree_modify(t2, 1);
        }
        if (t1 == 2)
        {
            link_modify(1, t2, 0);
        }
        if (t1 == 3)
        {
            cout << tree_query(t2) << endl;
        }
    }
}