[코드트리] 색깔 트리 (삼성 SW 역량테스트 2024 상반기 오후 2번, B형)
2022 하반기 이후 문제 풀이 시간이 3시간 → 4시간으로 변경,
A형 1문제 + B형 문제 1문제가 출제됩니다.
참고
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;
}