[Link Cut Tree] 4. 쿼리

이번 글에서는 트리에서의 경로 합을 구해봅시다.

일단 루트에서 노드 x까지의 쿼리를 구할 수 있다면, 두 노드 사이의 쿼리는 LCA를 이용해서 구할 수 있습니다.

( query(x, y) = query(x, root) + query(y, root) - query(lca(x, y), root) - query(parent(lca(x, y)), root) )

struct node {
    node *l, *r, *p;
    int cnt, val, sum;
    node() : l(0), r(0), p(0), cnt(0), val(0), sum(0) {}
};

일단 노드에 값과 구간 합을 저장하는 변수 val과 sum을 만들었습니다.

void update(node * x) {
    x->cnt = 1;
    x->sum = x->val;
    if (x->l) {
        x->cnt += x->l->cnt;
        x->sum += x->l->sum;
    }
    if (x->r) {
        x->cnt += x->r->cnt;
        x->sum += x->r->sum;
    }
}

update 함수도 sum을 구할 수 있도록 수정합니다.

이제 쿼리를 구하는 함수를 만들어 주는데, 두가지 방법이 있습니다.

query(x, y) = query(x, root) + query(y, root) - query(lca(x, y), root) - query(parent(lca(x, y)), root)를 이용합니다.

int query(node * x) {
    if (!x) return 0;
    access(x);
    return x->sum;
}

일단 x와 루트까지의 쿼리를 구하는 함수를 만듭니다.

x에 access하면 x와 root까지의 경로가 하나의 Splay Tree로 묶이게 되는데, 이 트리 전체의 sum이 x부터 root까지의 쿼리가 됩니다.

int query(node * x, node * y) {
    node * p = lca(x, y);
    return query(x) + query(y) - query(p) - query(parent(p));
}

x부터 y까지의 쿼리를 구하는 함수를 만듭니다.

이 방법은 최대값 쿼리와 같이 역연산이 불가능한 쿼리를 구할 수 없다는 문제점이 있습니다.

일단 x와 lca(x, y) 사이의 쿼리를 구하는 방법을 알아봅시다.

int query(node * x, node * y) {
    node * p = lca(y, x);
    int sum;
    access(x);
    splay(p);
    sum = p->val;
    if (p->r) sum += p->r->sum;
    
    access(y);
    splay(p);
    if (p->r) sum += p->r->sum;
    return sum;
}

x에 access하면 x와 root까지의 경로가 하나의 Splay Tree로 묶이게 되는데, 이 트리에서 LCA를 splay해주면 LCA에서 x까지의 경로가 LCA 노드의 오른쪽 자식에 묶이게 됩니다. (LCA는 제외)

이 오른쪽 자식의 sum값이 x에서 LCA 직전까지의 합이 됩니다.

LCA 노드의 값 + x에서 LCA 직전까지의 합 + y에서 LCA 직전까지의 합 = x에서 y까지의 합을 반환하면 됩니다.