본문 바로가기

자료구조

Persistent Segment Tree (PST)

$x, y$좌표가 $0$ 이상 $N$ 미만의 정수인 점 $N$개가 있다. $x, y$좌표는 모두 다르다. 이 때, 특정 직사각형 내부에 있는 점의 개수를 묻는 쿼리가 들어온다.

offline 풀이를 생각해보자. 직사각형 내부의 점의 개수를 구하는 쿼리는, $\sqsupset$ 형태의 구역에 있는 점들의 개수를 구하는 쿼리 2개로 쪼갤 수 있다. $\sqsupset$ 내부의 점의 개수를 구하는 쿼리들을을 $x$ 좌표($\sqsupset$에서 오른쪽 |) 순으로 정렬하고, 그 $x$ 좌표까지의 점들에 대해 $y$ 좌표를 BIT 또는 segment tree로 관리하면 $O(NlogN+QlogN)$에 해결할 수 있다.

이제 online 풀이를 생각해보자. offline과 같이 풀려고 하면 $O(N)$ 크기의 segment tree가 $N$개가 필요하다. 이렇게 하면 $O(N^2logN+NQlogN)$이다. (naive한 $O(NQ)$보다 안 좋아졌다.) 여기서 관찰할 수 있는 점은, 인접한 segment tree끼리는 변화가 적다는 것이다. 점 1개가 더 추가되었을 뿐이므로 $O(logN)$개의 노드들만 값이 바뀌게 된다. 이에 착안하여 만든 것이 Persistent Segment Tree이다. 주로 포인터를 이용해 구현하는데, 이미 만들어놓은 segment tree와 같은 부분은 새로 만들지 않고, 이미 있는 부분을 가리키게 하는 것이다. 그러면 이 문제를 $O(NlogN)$에 풀 수 있다.

PST는 트리에서도 사용할 수 있다. 한 정점과 루트와의 경로에 대한 segment tree를 만들어서 해결할 수 있는 문제들이 있다. 이 경우 부모에 대한 segment tree를 가져온 후 $O(logN)$개를 바꾸면 된다.

 

https://blog.myungwoo.kr/100  https://hongjun7.tistory.com/64 참고.

 

아래는 트리에서 간선들에 가중치가 있을 때, 어떤 두 정점을 잇는 경로에 있는 간선들의 가중치의 최댓값을 찾는 코드이다.

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

typedef pair<int, int> pii;

#define fi first
#define se second

const int MAXN = 50010;
const int INF = 1 << 30;

struct NOD {						// segment tree의 노드
	NOD *l, *r;					// 왼쪽과 오른쪽 자식 노드
	int cnt;					// 수의 개수
	NOD();
} *nil, *pst[MAXN];

NOD::NOD() : l(nil), r(nil), cnt(0) {}

vector<pii> ed[MAXN];

void mkpst(NOD *bf, NOD *nw, int l, int r, int x) {	// 이미 만들어놓은 bf에 x를 추가하여 nw를 만듦
	nw->cnt = bf->cnt + 1;
	if(l < r) {
		int m = (l + r) / 2;
		if(x <= m) {
			nw->r = bf->r;
			mkpst(bf->l, nw->l = new NOD(), l, m, x);
		}
		else {
			nw->l = bf->l;
			mkpst(bf->r, nw->r = new NOD(), m + 1, r, x);
		}
	}
}

int gpst(NOD *x, NOD *y, NOD *z, int l, int r) {	// 정점 x, y을 잇는 경로, lca(x, y) = z
	if(x->cnt + y->cnt - 2 * z->cnt == 0) return -1;
	if(l == r) return l;
	int m = (l + r) / 2;
	int t = gpst(x->r, y->r, z->r, m + 1, r);
	return ~t ? t : gpst(x->l, y->l, z->l, l, m);
}
int gpst(int x, int y, int z) { return gpst(pst[x], pst[y], pst[z], 0, INF); }

void dfs(int x, int p, int c) {				// dfs를 돌며 pst를 만든다
	pst[x] = new NOD();
	if(x != 1) mkpst(pst[p], pst[x], 0, INF, c);
	for(auto a : ed[x]) if(a.se != p) dfs(a.se, x, a.fi);
}

/*
new NOD()는 메모리를 새로 할당받는 것이 느리다
따라서 메모리를 미리 할당받아놓은 후에 newNod()라는 함수를 만들자
NOD nod[MX];
int nn;
NOD *newNOD() { return nod + nn++; }
*/

'자료구조' 카테고리의 다른 글

Segment tree beats (세그 비츠)  (0) 2022.02.28
Propogation 없는 lazy 세그트리  (0) 2021.10.29
Fenwick Tree (Binary Indexed Tree, BIT)  (0) 2021.04.09