이번 글에서는 트리에서의 경로 합을 구해봅시다.
일단 루트에서 노드 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까지의 쿼리를 구하는 함수를 만듭니다.
이 방법은 최대값 쿼리와 같이 역연산이 불가능한 쿼리를 구할 수 없다는 문제점이 있습니다.
- Splay Tree 이용
일단 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까지의 합을 반환하면 됩니다.
답글 남기기