백준 문제 풀이 - 사촌 9489번

📜 문제 이해하기

증가하는 정수 수열을 이용해서 트리를 만드는 방법은 다음과 같다.
첫 번째 정수는 트리의 루트 노드이다.
다음에 등장하는 연속된 수의 집합은 루트의 자식을 나타낸다. 이 집합에 포함되는 수의 첫 번째 수는 항상 루트 노드+1보다 크다.
그 다음부터는 모든 연속된 수의 집합은 아직 자식이 없는 노드의 자식이 된다. 그러한 노드가 여러 가지 인 경우에는 가장 작은 수를 가지는 노드의 자식이 된다.
집합은 수가 연속하지 않는 곳에서 구분된다.
예를 들어, 수열 1 3 4 5 8 9 15 30 31 32를 위의 규칙을 이용해 트리를 만들면 아래 그림과 같이 된다.
두 노드의 부모는 다르지만, 두 부모가 형제(sibling)일 때 두 노드를 사촌이라고 한다.
수열 특정 노드 번호 k가 주어졌을 때, k의 사촌의 수를 구하는 프로그램을 작성하시오.

💡 문제 재정의

증가하는 수열이 주어지고 연속된 수열이 자식 노드가 될 때 주어진 숫자의 사촌 노드의 갯수를 구하여라.

✏️ 계획 수립

1) Tree 구조를 만들고 데이터 삽입하기
먼저 Tree 구조를 만들어 Node를 삽입할 수 있도록 한다.
먼저 시작하는 root_num은 첫번재 숫자로 초기화한다. 그 후 parent가 될 숫자들을 모아놓은 parent_list에 root_num을 넣는다.
pointer를 선언하여 배열에서 어떤 부분을 가르킬지 정의한다.
child_list에 연속된 숫자를 가진 값들을 넣는다.
pointer의 값을 숫자의 개수만큼 갱신 시킨 뒤 parent_list에서 가장 작은 값을 꺼내 child_list에 속한 연속된 숫자들을 child로 만든다. 그 후 child_list를 parent_list에 접합시키고 num_list를 전부 순회할때까지 반복한다.

2) siblings 노드 찾기
다음의 노드를 siblings라고 정의할 수 있다.

• k와 같은 depth를 가진다.
• k와 부모가 같지 않다.
• k와 같은 조부모 노드를 가진다.
  위 조건을 만족하는 노드들의 갯수를 센 뒤 출력한다.

💻 계획 수행

import sys


class Node:
    def __init__(self, element, parent):
        self.element = element
        self.parent = parent
        self.children = []


class Tree:
    def __init__(self, root_element):
        self.root = Node(root_element, None)
        self.node_list = [self.root]

    def insert(self, parent_element, child_element):
        parent_node = self.find(parent_element)
        child_node = Node(child_element, parent_node)
        parent_node.children.append(child_node)
        self.node_list.append(child_node)

    def find(self, element):
        for node_ in self.node_list:
            if node_.element == element:
                return node_
        return None

    def get_depth(self, node_):
        depth = 0
        while node_.parent:
            depth += 1
            node_ = node_.parent
        return depth


def get_continuous_list(pointer_, num_list_):
    last_num = num_list_[pointer_]
    continuous_list = [last_num]

    for num in num_list_[pointer_ + 1:]:
        if num == last_num + 1:
            continuous_list.append(num)
            last_num = num
        else:
            break
    return continuous_list


if __name__ == '__main__':
    while True:
        n, k = map(int, sys.stdin.readline().split())
        if n == 0 and k == 0:
            break

        num_list = list(map(int, sys.stdin.readline().split()))

        root_num = num_list[0]
        tree = Tree(num_list[0])
        parent_list = [root_num]

        pointer = 1
        while pointer < n:
            child_list = get_continuous_list(pointer, num_list)
            pointer = pointer + len(child_list)
            parent = parent_list.pop(0)
            for child in child_list:
                tree.insert(parent, child)
            parent_list.extend(child_list)

        k_node = tree.find(k)
        k_depth = tree.get_depth(k_node)
        siblings = 0
        for node in tree.node_list:
            if k_depth == tree.get_depth(node) and node.parent != k_node.parent and node.parent.parent == k_node.parent.parent:
                siblings += 1

        print(siblings)

🤔 회고

Tree 자료구조를 사용하여 문제를 풀 수 있었다.
구현에 있어서 조금 시간이 걸렸지만 Tree를 적극적으로 사용해 볼 수 있어 좋은 경험이었던 것 같다.