비재귀 세그먼트 트리 - Efficient and easy segment tree

세그먼트 트리는 임의의 위치의 값들이 계속 변화하고, 특정 구간에 대한 연산(어떤 구간의 합, 어떤 구간 중 최소값 등)을 빠르게 구할 때 용이한 자료구조이다.

한국어로는 구간 합 트리?라고도 하는 것 같다.

이 글을 적는 이유는 세그먼트 트리 자체를 다루기 위한 것은 아니고, 크기를 2배로 잡는 대신에 비재귀 형태로 간단하게 작성할 수 있는 코드를 소개하려고 한다.

처음 접한 건 2016년 쯤 코드포스에서 읽었다. 워낙 간단하고 명료해서 아직까지 외우고 있다.

링크: http://codeforces.com/blog/entry/18051

대회에서도 자주 사용했는데, 블로그 포스팅은 한 적이 없어서 나중을 위해 다시 적어둔다.

자세한 설명은 원문 링크에 자세히 나와있으므로 생략한다.

사용할 때 주의할 점은, query는 [l, r)구간으로 사용하고 있으며 배열 크기가 N이라면 배열 인덱스의 시작은 N번부터라는 점이다.

0번 인덱스를 트리의 root로 사용하고 있기 때문에, 문제에서 일반적인 a[0], a[1], ... a[n-1] 은 모두 a[n], a[n+1], ... a[n+n-1] 에 저장하면 된다.

소스 코드

워낙 간단하기도 하지만, 너무 간단해서 연산이 이해가 안 가는 부분들이 있을 것이다.

t[i] = t[i<<1] + t[i<<1 | 1] 은 트리의 왼쪽 자식은 현재 위치 * 2, 오른쪽 자식은 현재 위치 * 2 + 1 을 이용한 것이다.
1을 OR 하는 이유는 짝수에 +1을 하는 것과 같다. 사실 + 1 로 해도 되지만 ,연산자 우선 순위를 따지면 저 표현이 괄호도 없이 깔끔하게 적을 수 있다.

반대로 update에서 t[pos] + t[pos ^ 1]은 짝수는 홀수로, 홀수는 짝수로 바꾸기때문에 부모/자식 접근을 반대로 하는 것이다.

query 함수에서 if(l & 1) 같은 것은 현재 보고있는 [l, r) 구간의 크기가 홀수이면 한 쪽의 원소가 누락된다. N=10 정도의 작은 케이스를 손으로 적어보면 이해가 잘 된다.