3차원 펜윅 트리(Fenwick Tree/Binary Indexed Tree/Segment Tree)

최근 3차원 공간의 부분합을 많이 구해야 일이 생겼다.

속도가 느려서 팬윅트리를 이용하여 해결하고자 했는데, 3차원으로 가니 직관적이지 않다.

조금 공부를 했고, 3차원 팬윅트리 코드를 완성 시켜보았다.

까먹지 않기 위해 자체 PS 문제를 제작, 솔루션을 남겨본다.

 

* 문제 설명

512^3 공간에서 인풋으로 주어진 범위 내의 '1' 을 카운트하여 반환하는 함수 작성

 

* 제약사항

힙 메모리 : 1gb 미만

스텍 메모리 : 1mb 미만

CPU TIME : 코드 참조

 

#include <ctime>
#include <cstdio>
#include <cstdlib>
#include <memory.h>

const int SIZE = 512;

// 코드를 읽고 PASS를 출력하게 하는 CountDefect 함수를 작성하시오.
// 제약사항 HEAP MEMORY < 1gb
int CountDefect(bool cube[SIZE][SIZE][SIZE], int xs, int xe, int ys, int ye, int zs, int ze)
{
	return 0;
}

bool cube[SIZE][SIZE][SIZE];
const int SEED = 3; // 변경될 수 있음.
const int MIN = 100;

int SlowCount(bool cube[SIZE][SIZE][SIZE], int xs, int xe, int ys, int ye, int zs, int ze)
{
	int res = 0;
	for (int z = zs; z <= ze; z++)
		for (int y = ys; y <= ye; y++)
			for (int x = xs; x <= xe; x++)
				if (cube[z][y][x])
					res++;
	return res;
}

int main(void)
{
	srand(SEED);
	memset(cube, 0, SIZE * SIZE * SIZE * sizeof(bool));
	for (int i = 0; i < 10000; i++) cube[rand() % SIZE][rand() % SIZE][rand() % SIZE] = 1;
	int xs, xe, ys, ye, zs, ze, t;
	clock_t start;
	long long score = 0;
	int elapse = 0;
	for (int T = 0; T < 1000; T++)
	{
		while (true)
		{
			xs = rand() % SIZE, xe = rand() % SIZE;
			ys = rand() % SIZE, ye = rand() % SIZE;
			zs = rand() % SIZE, ze = rand() % SIZE;
			if (xs > xe) t = xs, xs = xe, xe = t;
			if (ys > ye) t = ys, ys = ye, ye = t;
			if (zs > ze) t = zs, zs = ze, ze = t;
			if (xe - xs > MIN && ye - ys > MIN && ze - zs > MIN) break;
		}

		start = (int)clock();
		int ans = CountDefect(cube, xs, xe, ys, ye, zs, ze);
		elapse += (int)clock() - start;
		if (ans != SlowCount(cube, xs, xe, ys, ye, zs, ze)) score += 100000;
	}
	printf("%s\n", score + elapse < 10000 ? "PASS" : "FAIL");
	return 0;
}

 

 

솔루션

const int SIZE = 512;
bool init = 0;
const int SIZE_FW = SIZE + 1;
int fw[SIZE_FW][SIZE_FW][SIZE_FW];

int Sum(int x, int y, int z)
{
	int res = 0;
	for (int zz = z; zz > 0; zz -= (zz & -zz))
		for (int yy = y; yy > 0; yy -= (yy & -yy))
			for (int xx = x; xx > 0; xx -= (xx & -xx))
				res += fw[zz][yy][xx];
	return res;
}

int Sum(int xs, int xe, int ys, int ye, int zs, int ze)
{
	int res = Sum(xe, ye, ze);
	res -= Sum(xs - 1, ye, ze);
	res -= Sum(xe, ys - 1, ze);
	res -= Sum(xe, ye, zs - 1);
	res += Sum(xs - 1, ys - 1, ze);
	res += Sum(xs - 1, ye, zs - 1);
	res += Sum(xe, ys - 1, zs - 1);
	res -= Sum(xs - 1, ys - 1, zs - 1);
	return res;
}

void Add(int x, int y, int z)
{
	for (int zz = z; zz < SIZE_FW; zz += (zz & -zz))
		for (int yy = y; yy < SIZE_FW; yy += (yy & -yy))
			for (int xx = x; xx < SIZE_FW; xx += (xx & -xx))
				fw[zz][yy][xx]++;
}

void Init(bool cube[SIZE][SIZE][SIZE])
{
	init = 1;
	memset(fw, 0, SIZE_FW * SIZE_FW * SIZE_FW * sizeof(int));
	for (int z = 0; z < SIZE; z++)
		for (int y = 0; y < SIZE; y++)
			for (int x = 0; x < SIZE; x++)
				if(cube[z][y][x])
					Add(x + 1, y + 1, z + 1);
}

int CountDefect(bool cube[SIZE][SIZE][SIZE], int xs, int xe, int ys, int ye, int zs, int ze)
{
	if (!init) Init(cube);
	return Sum(xs+1, xe+1, ys+1, ye+1, zs+1, ze+1);
}

 

기존 시간 복잡도 : TC * SIZE^3

개선 시간 복잡도 : TC * (log(SIZE))^3