[코드트리] 산타의 선물공장 2

https://www.codetree.ai/frequent-problems/santa-gift-factory-2/description

코드트리

#include <cstdio>
#include <cstring>
#include <list>
#include <vector>
#include <iterator>


#define INIT (100)
#define MOVE (200)
#define SWAP (300) 
#define SEPERATE (400)
#define GET_INFO (500)
#define ABOUT_BELT (600)

int C;
int N, M;
std::list<int> Belt[100000 + 2];

std::list<int>::iterator pos[100000 + 2];
int prev[100000 + 2];
int next[100000 + 2];

void init_belt()
{
	// init
	for (int i = 0; i <= 100001; i++) {
		Belt[i].clear();
	}
	memset(pos, 0, sizeof(pos));

	// input
	scanf("%d %d", &N, &M);
	for (int b = 1; b <= M; b++) {
		int num = 0;
		scanf("%d", &num);
		
		// 앞뒤 조정
		if (Belt[num].empty()) {
			prev[b] = 0;
			next[b] = 0;
		}
		else {
			int node = Belt[num].back();
			prev[b] = node;
			next[b] = 0;
			next[node] = b;
		}
		std::list<int>::iterator iter = Belt[num].end();
		pos[b] = Belt[num].insert(iter, b);
	}
	return;
}


void move_all()
{
	int m_src = 0;
	int m_dst = 0;
	scanf("%d %d", &m_src, &m_dst);

	if (Belt[m_src].empty()) {
		printf("%d\n", Belt[m_dst].size());
		return;
	}

	// 앞뒤 조정
	int src_end = Belt[m_src].back();
	int dst_begin = Belt[m_dst].front();

	next[src_end] = dst_begin;
	prev[dst_begin] = src_end;


	std::list<int>::iterator iter = Belt[m_dst].begin();
	Belt[m_dst].splice(iter, Belt[m_src]);
	printf("%d\n", Belt[m_dst].size());
	return;
}


void swap_front()
{
	int m_src = 0;
	int m_dst = 0;
	scanf("%d %d", &m_src, &m_dst);


	// 둘다 물건 존재X
	if (Belt[m_src].empty() && Belt[m_dst].empty()) {
		printf("%d\n", Belt[m_dst].size());
		return;
	}
	// m_src에 물건 존재X	// dst -> src
	else if (Belt[m_src].empty() && !Belt[m_dst].empty()) {
		int box = Belt[m_dst].front();
		Belt[m_dst].pop_front();

		// 앞뒤 조정
		int dst_begin = Belt[m_dst].front();
		prev[dst_begin] = 0;
		prev[box] = 0;
		next[box] = 0;

		std::list<int>::iterator iter = Belt[m_src].begin();
		pos[box] = Belt[m_src].insert(iter, box);

		printf("%d\n", Belt[m_dst].size());
		return;
	}
	
	// m_dst에 물건 존재X	// src->dst
	else if (!Belt[m_src].empty() && Belt[m_dst].empty()) {
		int box = Belt[m_src].front();
		Belt[m_src].pop_front();


		// 앞뒤 조정
		int src_begin = Belt[m_src].front();
		prev[src_begin] = 0;
		prev[box] = 0;
		next[box] = 0;

		std::list<int>::iterator iter = Belt[m_dst].begin();
		pos[box] = Belt[m_dst].insert(iter, box);
		
		printf("%d\n", Belt[m_dst].size());
		return;
	}

	// 둘 다 물건 존재
	else if (!Belt[m_src].empty() && !Belt[m_dst].empty()) {
		int src2dst = Belt[m_src].front();
		int dst2src = Belt[m_dst].front();
		Belt[m_src].pop_front();
		Belt[m_dst].pop_front();

		// 앞뒤 조정
		int src_begin = Belt[m_src].front();
		int dst_begin = Belt[m_dst].front();
		next[src2dst] = dst_begin;
		next[dst2src] = src_begin;
		prev[src_begin] = dst2src;
		prev[dst_begin] = src2dst;

		std::list<int>::iterator iter2src = Belt[m_src].begin();
		std::list<int>::iterator iter2dst = Belt[m_dst].begin();
		pos[dst2src] = Belt[m_src].insert(iter2src, dst2src);
		pos[src2dst] = Belt[m_dst].insert(iter2dst, src2dst);
		
		printf("%d\n", Belt[m_dst].size());
		return;
	}
}


void seperate_half()
{
	int m_src = 0;
	int m_dst = 0;
	scanf("%d %d", &m_src, &m_dst);

	int n = Belt[m_src].size();

	if (Belt[m_src].empty() || n == 1) {
		printf("%d\n", Belt[m_dst].size());
		return;
	}


	std::list<int>::iterator src_iter = Belt[m_src].begin();
	std::advance(src_iter, n / 2);

	// 앞뒤 조정용
	int dst_begin = 0;
	if (!Belt[m_dst].empty()) dst_begin = Belt[m_dst].front();
	int move_end = 0;
	if (n / 2 == 1) move_end = Belt[m_src].front();
	else move_end = *(std::prev(src_iter, 1));

	Belt[m_dst].splice(Belt[m_dst].begin(), Belt[m_src], Belt[m_src].begin(), src_iter);

	// 앞뒤 조정
	int src_begin = Belt[m_src].front();
	prev[src_begin] = 0;
	if (dst_begin != 0) prev[dst_begin] = move_end;
	next[move_end] = dst_begin;

	//printf("dst : %d move: %d src :%d \n", dst_begin, move_end, src_begin);

	printf("%d\n", Belt[m_dst].size());
	return;
}


void get_box_info()
{
	int p_num = 0;
	scanf("%d", &p_num);

	std::list<int>::iterator iter = pos[p_num];

	int a = -1, b = -1;
	if (prev[*iter] != 0) a = prev[*iter];
	if (next[*iter] != 0) b = next[*iter];
	//printf("[prev]%d [next]%d\n", a, b);

	printf("%d\n", a + 2 * b);
	return;
}


void get_belt_info()
{
	int b_num = 0;
	scanf("%d", &b_num);

	int a = -1, b = -1, c = 0;
	if (!Belt[b_num].empty()) {
		a = Belt[b_num].front();
		b = Belt[b_num].back();
		c = Belt[b_num].size();
	}
	printf("%d\n", a + (2 * b) + (3 * c));
	return;
}


int main()
{
	scanf("%d", &C);
	for (int c = 0; c < C; c++) {
		int cmd = 0;
		scanf("%d", &cmd);
		switch (cmd)
		{
		case (INIT):
			init_belt();
			break;
		case (MOVE) :
			move_all();
			break;
		case (SWAP):
			swap_front();
			break;
		case (SEPERATE):
			seperate_half();
			break;
		case (GET_INFO):
			get_box_info();
			break;
		case (ABOUT_BELT):
			get_belt_info();
			break;
		}
	}
}

디버깅용 각 벨트에 있는 박스 확인 코드

void debug()
{

	for (int i = 1; i <= N; i++) {
		if (Belt[i].empty()) continue;
		printf("[%d]\n", i);
		for (int n : Belt[i]) {
			printf("(%d) %d (%d) ",(prev[n]), n, (next[n]));
		}
		printf("\n");
	}
	printf("======================\n");
}

1. 키포인트는 특정 박스의 앞과 뒤에 있는 박스 번호를 prev와 next라는 배열로 관리해 주는 것이다. 

그러면 벨트에 있는 박스 전체 혹은 일부를 옮길 때, 옮기는 박스 그룹의 가장 뒤에있는 박스와 남아있는 박스 그룹의 가장 앞에있는 박스의 각각 next, prev만 수정해주면 된다. 중간에 있는 박스들의 정보를 굳이 수정해주지 않아도 되어 훨씬 효율적인 것이다. 

 

2. 그리고 1번과 같이 관리해야 500번 함수를 구현하는 것이 쉬워진다.. 

만약 입력받은 숫자의 prev가 0이라면 앞에 아무것도 없는 것이므로 a = -1이다. 만약 next가 0이라면 뒤에 아무것도 없는 것이므로 b = -1 이다. 이러한 정보들을 모두 배열로 관리하고, 인덱스로 접근하기 때문에 시간 복잡도는 O(1) 일 것이다.