농장 관리 (백준 5916)

안녕하세요. 오늘은 농장을 관리할 거예요.

문제

https://www.acmicpc.net/problem/5916

아이디어

이 문제는 그냥 트리와 쿼리 1과 느리게 갱신되는 세그먼트 트리를 더해주면 됩니다.

소스코드

#include <iostream>
#include <vector>
#include <algorithm>
#define ll long long
using namespace std;
 
ll tree[404040] = { 0 }, lazy[404040] = { 0 };
void pp(ll s, ll e, ll node)
{
    if (lazy[node])
    {
        tree[node] += lazy[node];
        if (s != e)
        {
            lazy[node * 2] += lazy[node];
            lazy[node * 2 + 1] += lazy[node];
        }
        lazy[node] = 0;
    }
}
ll query(ll s, ll e, ll node, ll l, ll r)
{
    pp(s, e, node);
    if (e < l || r < s) return 0;
    if (l <= s && e <= r) return tree[node];
    ll mid = (s + e) / 2;
    return query(s, mid, node * 2, l, r) + query(mid + 1, e, node * 2 + 1, l, r);
}
void change(ll s, ll e, ll node, ll l, ll r)
{
    pp(s, e, node);
    if (e < l || r < s) return;
    if (l <= s && e <= r)
    {
        lazy[node]++;
        pp(s, e, node);
        return;
    }
    ll mid = (s + e) / 2;
    change(s, mid, node * 2, l, r);
    change(mid + 1, e, node * 2 + 1, l, r);
    tree[node] = tree[node * 2] + tree[node * 2 + 1];
}

ll in[101010] = { 0 }, out[101010] = { 0 }, par[101010] = { 0 }, dpt[101010] = { 0 }, sz[101010] = { 0 }, top[101010] = { 0 };
bool ck[101010] = { 0 };
vector <ll> down[101010];
vector <ll> graph[101010];

void dfs(ll node = 1)
{
    ck[node] = 1;
    for (ll next : graph[node])
    {
        if (ck[next]) continue;
        down[node].push_back(next);
        dfs(next);
    }
}
void dfs1(ll node = 1)
{
    sz[node] = 1;

    ll mx = 0, p = 0, idx = 0;
    for (ll next : down[node])
    {
        dpt[next] = dpt[node] + 1;
        par[next] = node;
        dfs1(next);
        sz[node] += sz[next];

        if (sz[next] > mx)
        {
            mx = sz[next];
            p = idx;
        }
        idx++;
    }
    if (down[node].size()) swap(down[node][0], down[node][p]);
}
ll pv = 0;
void dfs2(ll node = 1)
{
    in[node] = ++pv;
    for (ll next : down[node])
    {
        if (next == down[node][0])
            top[next] = top[node];
        else
            top[next] = next;
        dfs2(next);
    }
    out[node] = pv;
}

ll N;
void update(ll u, ll v)
{
    while (top[u] != top[v])
    {
        if (dpt[top[u]] < dpt[top[v]]) swap(u, v);
        ll st = top[u];
        change(1, N, 1, in[st], in[u]);
        u = par[st];
    }
    if (u == v) return;
    if (dpt[u] > dpt[v]) swap(u, v);
    change(1, N, 1, in[u] + 1, in[v]);
}
ll SUM(ll u, ll v)
{
    ll ans = 0;
    while (top[u] != top[v])
    {
        if (dpt[top[u]] < dpt[top[v]]) swap(u, v);
        ll st = top[u];
        ans += query(1, N, 1, in[st], in[u]);
        u = par[st];
    }
    if (u == v) return ans;
    if (dpt[u] > dpt[v]) swap(u, v);
    return ans + query(1, N, 1, in[u] + 1, in[v]);
}

int main(void)
{
    ios_base::sync_with_stdio(false); cin.tie(NULL);
    ll M, i, a, b;
    char c;

    cin >> N >> M;
    for (i = 0; i < N - 1; i++)
    {
        cin >> a >> b;
        graph[a].push_back(b);
        graph[b].push_back(a);
    }
    dfs(); dfs1(); dfs2();

    for (i = 0; i < M; i++)
    {
        cin >> c >> a >> b;
        if (c == 'P') update(a, b);
        else cout << SUM(a, b) << "\n";
    }
}

감사합니다.