2. 업그레이드 된 HEAP 힙 (특정원소 삭제, 업데이트 O(logN)만에 가능)

2023.03.12 - [자료구조/HEAP 힙] - 1.HEAP 힙 기본 구현 (C언어 구현, index 1부터 시작)

1.HEAP 힙 기본 구현 (C언어 구현, index 1부터 시작)

지난 시간에 구현한 힙 코드를 조금 더 업그레이드 시켜서 특정 원소 삭제 및 업데이트가 O(logN)만에 가능한 코드를 만들어 보겠습니다.

 

1. 기본 아이디어

지난번에는 특정 id에 해당하는 원소를 찾기위해 힙의 첫번째 원소부터 순차적으로 반복문을 돌리면서 해당하는 인덱스를 얻었습니다. 하지만, 이 과정을 없애버릴 수 있습니다. 따로 테이블에 특정 id에 해당하는 원소가 Heap의 몇 번 인덱스에 존재하는지를 기록하고 관리하면 됩니다. 이렇게 되면 O(N)이 걸렸던 것을 O(1)로 줄일 수 있습니다. 나머지 부분은 지난 시간의 아이디어와 동일합니다. 그래서 결과적으로 O(1+logN)==> O(logN)만에 특정 원소를 삭제하고 업데이트 할 수 있습니다.

 

2. 실제 구현 코드

저는 특정 id에 해당하는 원소가 Heap의 몇 번 인덱스에 존재하는지를 Unordered Map을 활용해서 기록했습니다. 참고로 id의 범위가 배열로 표현할 수 있을 만큼 작다면, 그냥 int형 배열을 쓰는게 당연히 더 효율적입니다. 그렇지만 실제 문제에서는 id의 범위가 "1억이하 자연수" 같이 크게 주어질 수 있는데 이 때는 Unordered Map을 쓰던지 아니면 그냥 Map을 쓰던지 해야합니다. 

 

실제로 구현할 때 특히 신경써야 하는 부분은 index를 갱신하는 것입니다. 새 원소가 push될 때 새 원소에게 index를 부여해야 하고, 원소가 pop될 때 그 원소의 index를 없애줘야 합니다. 또한 up,down연산에서 원소들이 swap될 때 index도 바뀌도록 처리해주는 걸 절대로 까먹어서는 안 됩니다.

#include<stdio.h>	//힙 응용(특정원소 삭제 O(logN), 업데이트 O(logN)
#include<unordered_map>
using namespace std;

const int MAX_N = 1000;
const int MAX_ID_NUM = 10000;

struct Student {
	int p1, p2, id;
};
Student heap[MAX_N];

int heapSize = 0;

unordered_map<int, int> indexMap;	//indexMap[학생 아이디]:힙에 몇번째 인덱스에 이 id 가진 학생이 존재하는지 저장

void heapInit() {
	heapSize = 0;
	indexMap.clear();
}

//p1큰게 앞에, 같다면 p2작은게 앞으로
int cmpStudent(int index1, int index2) {
	return (heap[index1].p1 > heap[index2].p1) || ((heap[index1].p1 == heap[index2].p1) && (heap[index1].p2 < heap[index2].p2));
}

void up(int index) {
	int cur = index;
	int parent = (cur >> 1);

	while (parent > 0 && cmpStudent(cur, parent)) {
		Student temp = heap[cur];
		heap[cur] = heap[parent];
		heap[parent] = temp;

		//인덱스 갱신
		indexMap[heap[cur].id] = cur;
		indexMap[heap[parent].id] = parent;

		cur = parent;
		parent = (cur >> 1);
	}
}

void down(int index) {
	int cur = index;
	int child = (cur << 1);

	while (child <= heapSize) {
		if (child + 1 <= heapSize) {
			child = cmpStudent(child, child + 1) ? child : child + 1;
		}
		if (cmpStudent(cur, child))break;
		Student temp = heap[cur];
		heap[cur] = heap[child];
		heap[child] = temp;

		//인덱스 갱신
		indexMap[heap[cur].id] = cur;
		indexMap[heap[child].id] = child;

		cur = child;
		child = (cur << 1);
	}
}

void push(Student data) {
	heap[++heapSize] = data;

	//새로 추가된 원소 인덱스 저장
	indexMap.insert(make_pair(data.id, heapSize));

	up(heapSize);
}

Student pop() {
	Student ret = heap[1];

	indexMap.erase(ret.id);	//삭제 원소는 인덱스 0으로 바꿔주자

	if (heapSize == 1) {
		heapSize--;
	}
	else {
		heap[1] = heap[heapSize--];
		indexMap[heap[1].id] = 1;	//인덱스 갱신
	}

	down(1);
	return ret;
}

//id에 해당하는 학생 힙에서 삭제하기
void remove(int id) {
	//1.해당 학생의 힙 인덱스 구하기
	int index = indexMap[id];

	//2.그 노드의 우선순위 최대로 만들어주기
	heap[index].p1 = 99999;	heap[index].p2 = -1;

	//3.루트로 올려주자
	up(index);

	//4. pop해주자
	pop();
}

void update(int id, int newP1, int newP2) {
	//1.해당 학생의 힙 인덱스 구하기
	int index = indexMap[id];

	//2.그 노드 우선순위 변화시키기
	heap[index].p1 = newP1;		heap[index].p2 = newP2;

	//3.위로 보내기
	up(index);

	//4. 아래로 보내기
	down(index);
}


int main(void) {

	heapInit();
	int p1 = 1;	int p2 = 2;	int id = 10;

	Student student;
	for (int i = 0; i < 10; i++) {
		student.p1 = p1;	student.p2 = p2;	student.id = id;

		push(student);

		if (i & 1)p1++;
		if (p2 == 1)p2 = 2;
		else p2--;
		id--;
	}
	//pop하면 이 순서임
	/*
	p1: 5, p2: 1, id: 1
	p1: 5, p2: 2, id: 2
	p1: 4, p2: 1, id: 3
	p1: 4, p2: 2, id: 4
	p1: 3, p2: 1, id: 5
	p1: 3, p2: 2, id: 6
	p1: 2, p2: 1, id: 7
	p1: 2, p2: 2, id: 8
	p1: 1, p2: 1, id: 9
	p1: 1, p2: 2, id: 10
	*/

	//id=6인 원소 삭제 해보자
	remove(6);
	/*
	p1: 5, p2: 1, id: 1
	p1: 5, p2: 2, id: 2
	p1: 4, p2: 1, id: 3
	p1: 4, p2: 2, id: 4
	p1: 3, p2: 1, id: 5
	p1: 2, p2: 1, id: 7
	p1: 2, p2: 2, id: 8
	p1: 1, p2: 1, id: 9
	p1: 1, p2: 2, id: 10
	*/

	//id=5인 원소 update 해보자.
	update(5, 10, 1);
	for (; heapSize != 0;) {
		Student ret = pop();
		printf("p1: %d, p2: %d, id: %d\n", ret.p1, ret.p2, ret.id);
	}
	printf("\n");

	/*
	p1: 10, p2 : 1, id : 5
	p1 : 5, p2 : 1, id : 1
	p1 : 5, p2 : 2, id : 2
	p1 : 4, p2 : 1, id : 3
	p1 : 4, p2 : 2, id : 4
	p1 : 2, p2 : 1, id : 7
	p1 : 2, p2 : 2, id : 8
	p1 : 1, p2 : 1, id : 9
	p1 : 1, p2 : 2, id : 10
	*/
}