백준 2927 남극탐험 Python

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

트리에서 노드 연결, 단일 업데이트, 구간합 쿼리를 수행하는 문제이다. 연결 여부는 union-find로 관리하고, 먼저 모든 연결 쿼리를 수행하여 HLD를 해준다. 이후, 순서대로 쿼리를 처리하면 된다. HLD 과정에서, 트리가 여러개로 분할된 경우도 존재하기 때문에 방문하지 않은 노드에 대해 모두 dfs를 실행해야 한다.

import sys
sys.setrecursionlimit(10**6)
rl = sys.stdin.readline
N = int(rl())
ns = list(map(int, rl().split()))
tree = [0] * (N*4)
 
def tree_update(node, l, r, x, v):
    if x < l or r < x: return tree[node]
    if l == r:
        tree[node] = v
        return tree[node]
    mid = (l+r) // 2
    tree[node*2] = tree_update(node*2, l, mid, x, v)
    tree[node*2+1] = tree_update(node*2+1, mid+1, r, x, v)
    tree[node] = tree[node*2] + tree[node*2+1]
    return tree[node]
 
def tree_query(node, l, r, s, e):
    if r < s or e < l: return 0
    if s <= l and r <= e: return tree[node]
    mid = (l+r) // 2
    return tree_query(node*2, l, mid, s, e) + tree_query(node*2+1, mid+1, r, s, e)
 
def dfs(x):
    vst[x] = 1
    for nx in es[x]:
        if vst[nx]: continue
        vst[nx] = 1
        g[x].append(nx)
        dfs(nx)
 
def dfs1(x):
    sz[x] = 1
    for i, nx in enumerate(g[x]):
        par[nx] = x
        dep[nx] = dep[x] + 1
        dfs1(nx)
        sz[x] += sz[nx]
        if sz[g[x][0]] < sz[g[x][i]]:
            g[x][0], g[x][i] = g[x][i], g[x][0]
 
pos = 0
def dfs2(x):
    global pos
    pos += 1
    inc[x] = pos
    for nx in g[x]:
        top[nx] = top[x] if nx == g[x][0] else nx
        dfs2(nx)
 
 
def query(a, b):
    ret = 0
    while top[a] != top[b]:
        if dep[top[a]] < dep[top[b]]: a, b = b, a
        at = top[a]
        ret += tree_query(1, 1, N, inc[at], inc[a])
        a = par[at]
    if dep[a] > dep[b]: a, b = b, a
    ret += tree_query(1, 1, N, inc[a], inc[b])
    return ret
 
 
def findp(x):
    if p[x] == x: return x
    p[x] = findp(p[x])
    return p[x]
 
def union(a, b):
    a = findp(a)
    b = findp(b)
    if a < b: a, b = b, a
    p[b] = a
 
Q = int(rl())
qs = []
es = [[] for _ in range(N+1)]
p = [i for i in range(N+1)]
vst = [0 for _ in range(N+1)]
g = [[] for _ in range(N+1)]
sz, dep, par, top, inc = [0]*(N+1), [0]*(N+1), [0]*(N+1), [0]*(N+1), [0]*(N+1)
 
for i in range(Q):
    q, a, b = rl().split()
    a, b = int(a), int(b)
    if q == "bridge":
        if findp(a) != findp(b):
            union(a, b)
            es[a].append(b)
            es[b].append(a)
        qs.append((q, a, b))
    else:
        qs.append((q, a, b))
        
for i in range(1, N+1):
    if not vst[i]:
        dfs(i); dfs1(i); dfs2(i)
 
for idx, v in enumerate(ns):
    tree_update(1, 1, N, inc[idx+1], v)
p = [i for i in range(N+1)]
for q, a, b in qs:
    if q == "bridge":
        if findp(a) != findp(b):
            union(a, b); print("yes")
        else:
            print("no")
    elif q == "penguins":
        tree_update(1, 1, N, inc[a], b)
    else:
        if findp(a) == findp(b):
            print(query(a, b))
        else:
            print("impossible")