코드트리 | 코딩테스트 준비를 위한 알고리즘 정석
국가대표가 만든 코딩 공부의 가이드북 코딩 왕초보부터 꿈의 직장 코테 합격까지, 국가대표가 엄선한 커리큘럼으로 준비해보세요.
www.codetree.ai
트리 문제다
트리 구현은 배열로 한다.
색깔 변경
특정 노드 를 루트로 하는 서브트리의
모든 노드의 색깔을 지정된 색 로 변경합니다.
최대 50000번 주어진다..
이걸 dfs로 구현하면 무조건 시간 초과나겠지.
점수 조회
깊이가 최대 20000인데
이걸 dfs를 다 한다는건
무조건 시간 초과 날 거 같다.
색이 변할 때마다
트리를 타고 올라가면서
색깔의 카운트를 바꾼다?
무조건 시간 초과 날 것이다..
방법은? "lazy propagation"이다!!
시간을 저장해서
색깔 조회나 점수 조회를 할 때만!!
lazy propagation을 해주자.
이런 테크닉은 삼성 B형에
자주 나온다.
바꾸라고 순진하게 바꾸는게 아니다 이 말.
1. 노드 추가
maxDepth 조건이 있다.
루트 노드이면 상관이 없고
루트 노드가 아니라면
재귀를 타고 루트까지 올라가며 -> canMakeChild() 함수
maxDepth 보다 크거나 같다면 false를 반환한다.
|
1
2
3
4
5
6
7
8
9
|
// 재귀로 루트까지 올라간다.
bool canMakeChild(Node cur, int need)
{
if (cur.id == 0)
return true;
if (cur.maxDepth <= need)
return false;
return canMakeChild(nodes[cur.parentId], need + 1);
}
|
cs |
2. 색깔 변경
노드의 색을 바꾸고 그 노드의 서브트리의 색을 전부 바꾼다.
-> naive 하게 하면 시간 초과가 날 게 뻔하다.
for (int i = 1; i <= q; i++)
for문으로 시간을 체크하고
nodes[mId].color = color;
nodes[mId].lastUpdate = i;
해당 노드의 색만 바꾸고
lastUpdate에 색깔 변경이 일어난 시간을 넣어준다.
3. 노드 색 출력
서브 트리의 노드의 색이 바뀌었을 수도 있으므로
이제 여기서 lastUpdate를 이용해서
lazy propagation을 해줘야한다.
우선 출력해야하는 노드에서
재귀를 통해 루트까지 올라간다.
거기서부터 lastUpdate를 비교하면서
색이 바뀌었는지를 체크한다.
그리고 그 색을 자식에게 넘긴다.
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
|
// first가 color second가 lastUpdate
pair<int, int> getColor(Node cur) {
// 루트의 부모는 id가 0
if (cur.id == 0)
return make_pair(0, 0);
// 재귀를 통해 루트까지 올라간다.
pair<int, int> info = getColor(nodes[cur.parentId]);
// 재귀를 탈출한다.
if (info.second > cur.lastUpdate)
return info;
// lastUpdate가 더 크다면
// cur의 색을 자식에게 보낸다.
else
return make_pair(cur.color, cur.lastUpdate);
}
|
cs |
4. 점수 조회
우선 color는 1~5까지의 값이다.
구조체와 구조체 함수를 선언하자.
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
|
struct ColorCount {
// false로 초기화를 해줘야한다
bool cnt[6] = { false, };
};
// 두 ColorCount 구조체를 합치는 함수
ColorCount addColorCount(const ColorCount& a, const ColorCount& b) {
ColorCount res;
for (int i = 1; i <= 5; i++) {
if(a.cnt[i] || b.cnt[i])
res.cnt[i] = true;
}
return res;
}
// ColorCount 구조체의 서로 다른 색 개수의 제곱을 반환하는 함수
int score(const ColorCount& a) {
int result = 0;
for (int i = 1; i <= 5; i++) {
if (a.cnt[i])
result += 1;
}
return result * result;
}
|
cs |
이건 루트에서 자식 노드로 재귀를 통해 이동하며
beauty에 더해가면 된다.
이 때 자식의 lastUpdate가 시간이 더 크다면
color를 바꾸어준다.
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
|
// 이건 자식으로 내려가야 한다
pair<int, ColorCount> getBeauty(Node cur, int color, int lastUpdate)
{
// lastUpdate 및 color 갱신
if (lastUpdate < cur.lastUpdate) {
lastUpdate = cur.lastUpdate;
color = cur.color;
}
pair<int, ColorCount> result;
result.second.cnt[color] = true;
for (int cId : cur.childIds)
{
Node child = nodes[cId];
pair<int,ColorCount> subResult = getBeauty(child, color, lastUpdate);
result.second = addColorCount(result.second, subResult.second);
result.first += subResult.first;
}
result.first += score(result.second);
return result;
}
|
cs |
<전체 코드>
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
|
#include <iostream>
#include <vector>
using namespace std;
struct Node {
int id;
int color;
int lastUpdate;
int maxDepth;
int parentId;
vector<int> childIds;
};
Node nodes[100005];
bool isRoot[100005];
vector<int> root_id;
// 재귀로 루트까지 올라간다.
bool canMakeChild(Node cur, int need)
{
if (cur.id == 0)
return true;
if (cur.maxDepth <= need)
return false;
return canMakeChild(nodes[cur.parentId], need + 1);
}
// first가 color second가 lastUpdate
pair<int, int> getColor(Node cur) {
// 루트의 부모는 id가 0
if (cur.id == 0)
return make_pair(0, 0);
// 재귀를 통해 루트까지 올라간다.
pair<int, int> info = getColor(nodes[cur.parentId]);
// 재귀를 탈출한다.
if (info.second > cur.lastUpdate)
return info;
// lastUpdate가 더 크다면
// cur의 색을 자식에게 보낸다.
else
return make_pair(cur.color, cur.lastUpdate);
}
struct ColorCount {
bool cnt[6] = { false, };
};
// 두 ColorCount 구조체를 합치는 함수
ColorCount addColorCount(const ColorCount& a, const ColorCount& b) {
ColorCount res;
for (int i = 1; i <= 5; i++) {
if(a.cnt[i] || b.cnt[i])
res.cnt[i] = true;
}
return res;
}
// ColorCount 구조체의 서로 다른 색 개수의 제곱을 반환하는 함수
int score(const ColorCount& a) {
int result = 0;
for (int i = 1; i <= 5; i++) {
if (a.cnt[i])
result += 1;
}
return result * result;
}
// 이건 자식으로 내려가야 한다
pair<int, ColorCount> getBeauty(Node cur, int color, int lastUpdate)
{
// lastUpdate 및 color 갱신
if (lastUpdate < cur.lastUpdate) {
lastUpdate = cur.lastUpdate;
color = cur.color;
}
pair<int, ColorCount> result;
result.second.cnt[color] = true;
for (int cId : cur.childIds)
{
Node child = nodes[cId];
pair<int,ColorCount> subResult = getBeauty(child, color, lastUpdate);
result.second = addColorCount(result.second, subResult.second);
result.first += subResult.first;
}
result.first += score(result.second);
return result;
}
int main() {
int q;
cin >> q;
// lastUpdate 시간이 필요하므로
// 1부터 시작하는 for문을 이용
for (int i = 1; i <= q; i++) {
int t, mId, pId, color, maxDepth;
cin >> t;
if (t == 100)
{
cin >> mId >> pId >> color >> maxDepth;
if (pId == -1)
{
root_id.push_back(mId);
isRoot[mId] = true;
}
if (isRoot[mId] || canMakeChild(nodes[pId], 1)) {
nodes[mId].id = mId;
nodes[mId].color = color;
nodes[mId].maxDepth = maxDepth;
nodes[mId].parentId = isRoot[mId] ? 0 : pId;
nodes[mId].lastUpdate = i;
if (!isRoot[mId])
nodes[pId].childIds.push_back(mId);
}
}
else if (t == 200)
{
cin >> mId >> color;
nodes[mId].color = color;
nodes[mId].lastUpdate = i;
}
else if (t == 300)
{
cin >> mId;
cout << getColor(nodes[mId]).first << '\n';
}
else if (t == 400)
{
int beauty = 0;
for (auto k : root_id)
{
if(isRoot[k])
beauty += getBeauty(nodes[k], nodes[k].color, nodes[k].lastUpdate).first;
}
cout << beauty << '\n';
}
}
}
|
cs |
'트리' 카테고리의 다른 글
| Directory (0) | 2024.01.28 |
|---|---|
| 최소 공통 조상 (LCA) (2) | 2024.01.13 |
| 리프 노드 개수 구하기 (0) | 2024.01.12 |
| 트리의 부모 찾기 (0) | 2024.01.12 |
| 트리의 기본 (0) | 2023.03.16 |