Index Tree (구간 합 구하기)

구간 합 문제

알고리즘 문제를 풀다보면 특정 구간에 대해서 합, 최대최소 등을 구하는 문제가 자주 등장한다.

예제 : 구간 합 구하기 4

위 문제는 포인터 두 개를 사용하여 적절한 구간을 매번 선택하면서 끝까지 탐색시키는 방식으로 간단하게 풀 수 있다.

#include <iostream>

using namespace std;

int arr[100001];

int		main(void)
{
	int		n, m;
	int		i;
	int		a, b;
	
	i = 0;
	scanf("%d %d", &n, &m);
	while (i++ < n)
	{
		scanf("%d", &arr[i]);
		arr[i] = arr[i] + arr[i - 1];
	}
	while (m--)
	{
		scanf("%d %d", &a, &b);
		printf("%d\n", arr[b] - arr[a - 1]);
	}
	return (0);
}

우선 Array에 문제에서 사용할 값들을 입력받는다.
입력받음과 동시에 직전 index에 저장되어 있던 값을 더해 현재 입력받은 수를 갱신한다.
이렇게 데이터를 입력받으면 각 index의 수들은 0번 index부터 자기 자신까지의 누적합을 담게 된다.

위 형태와 같이 누적합 배열을 만들었다면, N부터 M까지의 구간합은

Array[M] - Array[N - 1]

로 쉽게 계산된다는 것을 알 수 있다.

데이터가 계속 갱신되는 구간 합 문제

이제 문제를 확장해서 값이 계속 갱신되는 상태에서의 구간합을 구하는 방법에 대해 생각해보자.

만약 4번 index의 숫자를 10으로 갱신한다고 가정하면

4번 index부터 데이터의 size까지 누적합을 전부 다시 계산해주어야 한다.
위의 그림에서는 데이터가 10개 뿐이지만, 만약 데이터의 개수가 10만개, 100만개로 늘어난다면 선형시간 안에 누적합을 계산하지 못하게 될 것이다.

이를 해결하기 위한 알고리즘으로는 Segment Tree, Fenwick Tree, Index Tree 등이 있는데 이번 글에서 다루고자 하는 것은 Bottom Up 방식으로 데이터를 갱신하는 Index Tree이다.

Index Tree

예제 : 구간 합 구하기

Index Tree 만들기

Index Tree의 초기 모습은 아래와 같다.

데이터를 입력받으면 위 데이터를 이렇게 이진트리의 가장 말단에 순서대로 넣는다.
실제 데이터는 이러한 트리 형태가 아닌 1차원 배열의 형태로 저장된다는 것을 기억하자.

잘 와닿지 않는다면 아래 그림을 참고하면 된다.

이러한 형태로 데이터를 저장하기 위해서는 데이터가 모두 저장될 수 있는 이진트리의 깊이를 계산한 뒤, 해당하는 깊이의 start_index부터 값을 저장해야 한다.

시작 index 찾기

위 문제에서는 데이터의 개수가 총 5개인데, 이진트리는 깊이가 깊어질수록 노드의 개수는 2의 거듭제곱 형태로 증가하게 된다. 따라서 2의 거듭제곱중 5보다 크면서 가장 작은 값이 start_index가 된다.

입력으로 받은 숫자의 개수 5를 비트로 표현하면 101이다.
그리고 구해야하는 start_index는 2의 세제곱인 8이다.
즉, 들어온 수의 개수를 2진수로 표현했을 때 가장 큰 값을 표현하는 비트가 몇 번째 비트인지를 이용하면 start_index를 구할 수 있다.

int	find_start_index(int n)
{
	int	start_index;

	start_index = 1;
	while (n)
	{
		n >>= 1;
		start_index <<= 1;
	}
	return (start_index);
}

start_index를 구했다면, start_index부터 입력받은 값을 저장하면 된다.

Index Tree 값 채우기

위 그림에서 부모 노드의 값은 자식 노드들의 값의 합이다.
이들의 index 관계를 일반화하면 아래와 같다.

• leftchild_index == parents_index * 2
• rightchild_index == parents_index * 2 + 1

따라서 값을 갱신하는 과정을 코드로 작성하면 아래와 같다.

void	init_idt(int start_index)
{
	int		i;
	
	i = start_index - 1;
	while (0 < i)
	{
		arr[i] = arr[i * 2] + arr[i * 2 + 1];
		i--;
	}
}

이 과정이 끝나면 아래와 같은 형태가 된다.

값 update

만약 문제에서 제시된 것 처럼 3번 index의 값을 6으로 갱신하고자 한다면 내가 최종적으로 갱신해야하는 노드의 개수는 언제나 트리의 높이와 같다.

따라서 M개의 값을 갱신하는데 O(Mlog N) (M == 값 갱신 횟수) (N == start_index * 2)의 시간이 보장된다고 볼 수 있다.

갱신해야하는 값들을 색칠하면 위와 같은 형태이다.

이를 구현하기 위해서는 부모 노드를 찾는 방법을 생각해야한다.
이는 아까 자식노드의 index를 찾을 때 일반화했던 식의 역이므로

parents_index = child_index / 2

위와 같이 간단하게 정의할 수 있다.

따라서 이를 코드로 옮기면

void	update(int start_index, int index, long long num)
{
	int	i;
	
	i = start_index - 1 + index;
	arr[i] = num;
	while (0 < i)
	{
		i >>= 1;
		if (i != 0)
			arr[i] = arr[i * 2] + arr[i * 2 + 1]; 
	}
}

Index Tree로 구간 합 구하기

이제 Index Tree를 이용해서 어떻게 구간합을 구할 수 있을지 생각해보자.

만약 문제에서 제시한 것 처럼 2 ~ 5번 index의 구간합을 구하고자 하면

위 구간의 구간합을 구해야한다.
이 때 3~4번 index에 위치하는 값들의 합은 5번 index의 노드가 가지고 있으므로 Index Tree에서 구간합은

위 값들의 합을 구하면 된다.

이를 코드로 구현하기 위해 몇가지 case를 생각해보자.

1. 구간합을 구해야 하는 range의 left index가 짝수인 경우

left index가 짝수라는 의미는 자기 자신이 부모 노드의 left child라는 의미와 동일하다. 따라서 자기 자신을 구간 합 결과에 더하지 않더라도 부모 노드에 자기 자신의 값이 포함되어 있기 때문에 부모노드로 그냥 올라가면 된다.

2. 구간합을 구해야하는 range의 left index가 홀수인 경우

left index가 홀수라는 의미는 자기 자신이 부모 노드의 right child라는 의미와 동일하다. 즉, 부모 노드로 올라가면 범위를 벗어난 값을 포함하고 있을 것이므로, 자기 자신을 구간 합 결과에 더한 뒤 부모 노드의 오른쪽에 위치한 노드로 올라가야 한다.

3. 구간합을 구해야 하는 range의 right index가 짝수인 경우

right index가 홀수라는 의미는 자기 자신이 부모 노드의 left child라는 의미와 동일하다. 즉, 부모 노드로 올라가면 범위를 벗어난 값을 포함하고 있을 것이므로, 자기 자신을 구간 합 결과에 더한 뒤 부모 노드의 왼쪽에 위치한 노드로 올라가야 한다.

4. 구간합을 구해야하는 range의 right index가 홀수인 경우

right index가 짝수라는 의미는 자기 자신이 부모 노드의 right child라는 의미와 동일하다. 따라서 자기 자신을 구간 합 결과에 더하지 않더라도 부모 노드에 자기 자신의 값이 포함되어 있기 때문에 부모노드로 그냥 올라가면 된다.

종료 조건?

만약 위와 같이 올라가다가 left index가 right index보다 커졌을 경우 연산을 멈추고 종료하면 된다.

자세한 것은 코드를 보고 파악하자.

long long	get_sum(int start_index, int left, int right)
{
	long long	ans;
	
	ans = 0;
	left += start_index - 1;
	right += start_index - 1;
	while (left <= right)
	{
		if (left % 2 == 1)
			ans = ans + arr[left];
		if (right % 2 == 0)
			ans = ans + arr[right];
		left = (left + 1) / 2;
		right = (right - 1) / 2;
	}
	return (ans);
}

left와 right index를 갱신할 때 각각 +1, -1을 해 준 뒤 2로 나누는 이유는 두 index가 모이는 방향으로 올라도록 하기 위한 것이다.

Index Tree의 확장

Index Tree의 가장 큰 특징은 구간합 뿐만 아니라 각종 대푯값들에 대한 정보도 담을 수 있다는 것이다.

예를 들자면, 구간 최솟값을 결정할 때에도 사용할 수 있다.
최솟값

형태는 똑같다. 특정 범위가 주어지면 Index Tree 구간합을 구하는 것 대신 범위들의 대표값을 이용해 최솟값 연산을 하는 것으로 바꾸어 주면 된다.

즉, 구간이 주어지면 실제 비교해야하는 값은 위 진한 주황색으로 색칠된 두 값 뿐이라는 것이다.

이를 활용하면 구간 최댓값이나, 구간 GCD와 같은 문제도 조금씩만 변형하면 해결할 수 있게 된다.