2250번 : 트리의 높이와 너비

문제

 

---

풀이

 

루트노드를 직접 찾아야 한다는 것을 몰라서 많이 헤맸다. 그것만 아니었어도 30분은 아꼈을텐데...

이 문제를 보자마자 제일 먼저 한 생각은 레벨 순회를 해야겠다는 것이었다.

dfs에 속한다고 볼 수 있는 preorder, inorder, postorder와 달리 level traversal은 bfs에 속한다.

아무튼 이걸 사용해서 각 레벨별 정점을 정리해야겠다고 생각했다.

 

level1 - 1

level2 - 2, 3

level3 - 4, 5, 6, 7

 

이런식으로 저장될테니 나중에 레벨마다의 너비를 구할때도 각 레벨의 첫번째 값과 마지막 값의 길이만 구하면 된다.

 

근데 문제는 각 정점의 열을 구하는 것이었다...

처음에는 분할정복으로 각 정점의 열을 구했는데 이것도 문제는 없는 코드였다. 굳이 뭐 짧게 보여주자면

int cntNodes(int node) {
    if (node == -1)
        return 0;
    int cnt = cntNodes(tree[node].first) + cntNodes(tree[node].second) + 1;
    return cnt;
}

void findPos(int node, int left) {
    if (node == -1)
        return;
    int pos = left + cntNodes(tree[node].first);
    position[node] = pos;
    findPos(tree[node].first, left);
    findPos(tree[node].second, pos + 1);
}

이런식으로...구현했는데

cntNodes 함수는 각 정점과 그 정점의 자식 정점들의 합을 구하는 함수였다.

위 그림에서 1번 정점의 경우라면 7을 리턴했을 것이다.

 

findPos 함수에서는 이 cntNodes 함수를 이용해 열좌표를 구하고 왼쪽 오른쪽으로 분할했다...

근데 이것보다 훨씬 좋은 방법이 있었다.

 

내가 사용한 방법은 (왼쪽 트리)(루트)(오른쪽 트리)로 분할한 것과 같다.

근데 이건 inorder, 중위 순회와 같은 방식이다.

그니까 그냥 중위 순회하면 굳이 저런 복잡한 함수 쓸 것 없이 열을 구할 수 있는 것이다.

 

아무튼 각 정점의 level과 열 정보를 알았으니 나머지는 간단하다.

---

소스코드

 

#include <iostream>
#include <queue>
#include <map>
#include <vector>
#include <utility>

using namespace std;

map<int, pair<int, int>> tree; //트리
vector<vector<int>> levels; //레벨별로 각 정점 저장
vector<int> row; //각 정점의 열

int levelTraversal(int root) { //레벨 순회
    int cur, level;
    queue<pair<int, int>> node;
    node.push(make_pair(root, 1));
    while (!node.empty()) {
        cur = node.front().first;
        level = node.front().second;
        node.pop();
        levels[level].push_back(cur);
        if (tree[cur].first != -1)
            node.push(make_pair(tree[cur].first, level + 1));
        if (tree[cur].second != -1)
            node.push(make_pair(tree[cur].second, level + 1));
    }
    return level; //높이 반환
}

int cnt = 1;
void inorder(int node) { //중위 순위하며 각 정점의 열 저장
    if (tree[node].first != -1)
        inorder(tree[node].first);
    row[node] = cnt++;
    if (tree[node].second != -1)
        inorder(tree[node].second);
}

int main() {
    int N, node, l, r, root;
    vector<bool> parent; //root 확인용
    pair<int, int> result = {1, 1}; //최대 너비, 레벨

    cin >> N;
    levels.assign(N + 1, vector<int>(0));
    row.assign(N + 1, 0);
    parent.assign(N + 2, false);
    for (int i = 0; i < N; i++) {
        cin >> node >> l >> r;
        tree[node] = make_pair(l, r);
        parent[l + 1] = parent[r + 1] = true;
    }
    for (int i = 1; i <= N + 1; i++) { //부모 정점이 없다면 root
        if (!parent[i])
            root = i - 1;
    }

    int height = levelTraversal(root); //레벨 순회
    inorder(root); //중위 순회
    for (int i = 1; i <= height; i++) { //각 레벨마다
        int width = row[levels[i][levels[i].size() - 1]] - row[levels[i][0]] + 1; //(마지막 값-첫번째 값) = 해당 레벨의 너비
        if (width > result.second) { //갱신
            result.first = i;
            result.second = width;
        }
    }
    cout << result.first << ' ' << result.second;
}

전체 소스코드다.

 

map<int, pair<int, int>> tree; //트리
vector<vector<int>> levels; //레벨별로 각 정점 저장
vector<int> row; //각 정점의 열

트리를 저장할 map형 tree와 각 레벨별 정점을 저장할 levels 2차원 벡터다.

levels[2]에는 level이 2인 정점들이 저장될 것이다.

그리고 각 정점의 열을 저장할 row 벡터이다. row[3]에는 3번 정점의 열이 저장되는거고...

 

    int N, node, l, r, root;
    vector<bool> parent; //root 확인용
    pair<int, int> result = {1, 1}; //최대 너비, 레벨

    cin >> N;
    levels.assign(N + 1, vector<int>(0));
    row.assign(N + 1, 0);
    parent.assign(N + 2, false);
    for (int i = 0; i < N; i++) {
        cin >> node >> l >> r;
        tree[node] = make_pair(l, r);
        parent[l + 1] = parent[r + 1] = true;
    }
    for (int i = 1; i <= N + 1; i++) { //부모 정점이 없다면 root
        if (!parent[i])
            root = i - 1;
    }

입력을 처리하는 부분이다. 입력을 받으면서 부모 정점의 존재 여부를 체크한다.

parent[l]이 아니라 parent[l+1]인 이유는 -1이 들어올 수 있기 때문이다.

이렇게 1씩 증가해서 체크했으니 root를 체크할 땐 마지막에 1을 빼준다.

 

    int height = levelTraversal(root); //레벨 순회

이렇게 구한 root에 대해 레벨 순회를 하며 levels벡터를 채움과 동시에 트리의 전체 높이를 리턴받는다.

 

int levelTraversal(int root) { //레벨 순회
    int cur, level;
    queue<pair<int, int>> node;
    node.push(make_pair(root, 1));
    while (!node.empty()) {
        cur = node.front().first;
        level = node.front().second;
        node.pop();
        levels[level].push_back(cur);
        if (tree[cur].first != -1)
            node.push(make_pair(tree[cur].first, level + 1));
        if (tree[cur].second != -1)
            node.push(make_pair(tree[cur].second, level + 1));
    }
    return level; //높이 반환
}

레벨 순회를 하는 함수는 bfs와 거의 같다.

 

    inorder(root); //중위 순회

그리고 중위 순회를 하며 각 정점의 열을 저장한다.

 

int cnt = 1;
void inorder(int node) { //중위 순위하며 각 정점의 열 저장
    if (tree[node].first != -1)
        inorder(tree[node].first);
    row[node] = cnt++;
    if (tree[node].second != -1)
        inorder(tree[node].second);
}

이것도 그냥 dfs랑 유사하다...

 

    for (int i = 1; i <= height; i++) { //각 레벨마다
        int width = row[levels[i][levels[i].size() - 1]] - row[levels[i][0]] + 1; //(마지막 값-첫번째 값) = 해당 레벨의 너비
        if (width > result.second) { //갱신
            result.first = i;
            result.second = width;
        }
    }

각 레벨의 너비는 마지막 값의 열 좌표에서 첫번째 값의 열 좌표를 빼고 거기에 1을 더한 값이다.

최대 너비와 비교하며 갱신한다.