[코드트리] 색깔 트리 (삼성 SW 역량테스트 2024 상반기 오후 2번, B형)

삼성 A형 전체 링크

삼성 B형 전체 링크

 

2022 하반기 이후 문제 풀이 시간이 3시간 → 4시간으로 변경,

A형 1문제 + B형 문제 1문제가 출제됩니다.

 

참고

- BOJ 1707 : 이분 그래프 (vector)

 

https://www.codetree.ai/training-field/frequent-problems/problems/color-tree

이후 설명 및 입/출력은 링크 참고

 

노드 추가

 

NODE는 다음과 같이 정의한다. 

자식 노드의 개수의 제한이 없기 때문에 vector를 사용하였다.

각 NODE가 root인지 체크하는 checkRoot도 선언한다.

typedef struct st
{
	int id;
	int color;
	int maxDepth;
	int parent;
	vector<int> child;
}NODE;

NODE node[MAX];
int checkRoot[MAX];

 

p_id가 -1인 경우는 루트 노드가 된다.

 

루트 노드이거나,

주어진 노드의 모든 부모 노드에 대해 depth를 검사해서 child를 추가할 수 있는지 확인(checkMakeChild)하고,

노드를 추가하고, 부모 노드가 아닌 경우는 부모 노드의 자식에 노드를 추가한다.

void addNode(int m_id, int p_id, int color, int max_depth)
{
	if (p_id == -1) checkRoot[m_id] = PARENT;

	if (checkRoot[m_id] == PARENT || checkMakeChild(node[p_id], 1))
	{
		node[m_id].id = m_id;
		node[m_id].color = color;
		node[m_id].maxDepth = max_depth;
		node[m_id].parent = checkRoot[m_id] == PARENT ? 0 : p_id;

		if (checkRoot[m_id] != PARENT)
			node[p_id].child.push_back(m_id);
	}
}

 

checkMakeChild는 다음과 같다.

depth를 늘려가면서 모든 부모에 대해 maxDepth가 초과하지 않는지 확인한다.

int checkMakeChild(NODE cur, int depth)
{
	if (cur.id == 0) return 1; // 최대 깊이 이전에 부모를 찾은 경우
	if (cur.maxDepth <= depth) return 0;
	return checkMakeChild(node[cur.parent], depth + 1);
}

---

색깔 변경

 

주어진 노드와 자식 노드의 색깔을 dfs를 이용해 변경한다.

void dfs(int id, int color)
{
	node[id].color = color;

	for (int childId : node[id].child)
		dfs(childId, color);
}

void changeColor(int m_id, int color)
{
	dfs(m_id, color);
}

---

색깔 조회

 

주어진 노드의 색깔을 출력한다.

void retrieveColor(int m_id)
{
	printf("%d\n", node[m_id].color);
}

---

점수 조회

 

모든 루트 노드에 대해 점수를 계산하여 합한 값을 출력한다.

int getScore()
{
	int score = 0;
	int colorCount[MAX_COLOR + 1] = { 0 };
	for (int i = 1; i <= 100000; i++)
		if (checkRoot[i] == PARENT)
			score += calculate(node[i], node[i].color, colorCount);

	return score;
}

void retrieveScore()
{
	int score = getScore();

	printf("%d\n", score);
}

 

점수도 dfs를 이용해 계산한다.

colorCount 배열에 주어진 노드의 color를 1로 표시한다.

그리고 해당 노드의 모든 자식 노드에 대해 color를 찾는다.

자식의 colorCount를 tempColorCount에 더하고, 점수도 누적한다.

count에 tempColorCount에서 1 이상의 값인 경우 1을 더한다 (!! 이용)

count의 제곱을 sum에 누적한다.

int calculate(NODE cur, int colorCount[])
{
	int tempColorCount[MAX_COLOR + 1] = { 0 };
	tempColorCount[cur.color] = 1;

	int sum = 0;
	for (int childId : cur.child)
	{
		NODE child = node[childId];

		int childColorCount[MAX_COLOR + 1] = { 0 };
		int score = calculate(child, childColorCount);

		for (int i = 1; i <= MAX_COLOR; i++) tempColorCount[i] += childColorCount[i];

		sum += score;
	}

	int count = 0;
	for (int i = 1; i <= 5; i++) count += !!tempColorCount[i];

	sum += (count * count);

	for (int i = 1; i <= MAX_COLOR; i++) colorCount[i] += tempColorCount[i];

	return sum;
}

---

전체 코드는 다음과 같다.

#include <stdio.h>
#include <vector>

using namespace std;

#define MAX (100000 + 5000)
#define MAX_DEPTH (100 + 10)
#define MAX_COLOR (5)

#define ADD_NODE (100)
#define CHANGE_COLOR (200)
#define RETRIEVE_COLOR (300)
#define RETRIEVE_SCORE (400)

#define PARENT (1)

typedef struct st
{
	int id;
	int color;
	int maxDepth;
	int parent;
	vector<int> child;
}NODE;

NODE node[MAX];
int checkRoot[MAX];

int checkMakeChild(NODE cur, int depth)
{
	if (cur.id == 0) return 1; // 최대 깊이 이전에 부모를 찾은 경우
	if (cur.maxDepth <= depth) return 0;
	return checkMakeChild(node[cur.parent], depth + 1);
}

void addNode(int m_id, int p_id, int color, int max_depth)
{
	if (p_id == -1) checkRoot[m_id] = PARENT;

	if (checkRoot[m_id] == PARENT || checkMakeChild(node[p_id], 1))
	{
		node[m_id].id = m_id;
		node[m_id].color = color;
		node[m_id].maxDepth = max_depth;
		node[m_id].parent = checkRoot[m_id] == PARENT ? 0 : p_id;

		if (checkRoot[m_id] != PARENT)
			node[p_id].child.push_back(m_id);
	}
}

void dfs(int id, int color)
{
	node[id].color = color;

	for (int childId : node[id].child)
		dfs(childId, color);
}

void changeColor(int m_id, int color)
{
	dfs(m_id, color);
}

void retrieveColor(int m_id)
{
	printf("%d\n", node[m_id].color);
}

int calculate(NODE cur, int colorCount[])
{
	int tempColorCount[MAX_COLOR + 1] = { 0 };
	tempColorCount[cur.color] = 1;

	int sum = 0;
	for (int childId : cur.child)
	{
		NODE child = node[childId];

		int childColorCount[MAX_COLOR + 1] = { 0 };
		int score = calculate(child, childColorCount);

		for (int i = 1; i <= MAX_COLOR; i++) tempColorCount[i] += childColorCount[i];

		sum += score;
	}

	int count = 0;
	for (int i = 1; i <= 5; i++) count += !!tempColorCount[i];

	sum += (count * count);

	for (int i = 1; i <= MAX_COLOR; i++) colorCount[i] += tempColorCount[i];

	return sum;
}

int getScore()
{
	int score = 0;
	int colorCount[MAX_COLOR + 1] = { 0 };
	for (int i = 1; i <= 100000; i++)
		if (checkRoot[i] == PARENT)
			score += calculate(node[i], colorCount);

	return score;
}

void retrieveScore()
{
	int score = getScore();

	printf("%d\n", score);
}

int main()
{
	int Q;

	scanf("%d", &Q);

	for (int q = 0; q < Q; q++)
	{
		int COMMAND;

		scanf("%d", &COMMAND);

		if (COMMAND == ADD_NODE)
		{
			int m_id, p_id, color, max_depth;
			scanf("%d %d %d %d", &m_id, &p_id, &color, &max_depth);
			addNode(m_id, p_id, color, max_depth);
		}
		else if (COMMAND == CHANGE_COLOR)
		{
			int m_id, color;
			scanf("%d %d", &m_id, &color);
			changeColor(m_id, color);
		}
		else if (COMMAND == RETRIEVE_COLOR)
		{
			int m_id;
			scanf("%d", &m_id);
			retrieveColor(m_id);
		}
		else if (COMMAND == RETRIEVE_SCORE)
		{
			retrieveScore();
		}
	}

	return 0;
}