[코드트리][Python] 삼성 SW 역량테스트 2024 상반기 오후 2번 문제 - 색깔 트리

📚 문제

https://www.codetree.ai/training-field/frequent-problems/problems/color-tree/description?page=1&pageSize=5

코드트리 | 코딩테스트 준비를 위한 알고리즘 정석

 

삼성전자 코테를 앞두고 기출을 풀어보고자 준비했다. 실제 상반기에 풀었던 문제 그대로였다. 

(따로 글은 작성하지 않았지만 코테 통과 후 면접에서 해당 사업부에서는 신입으로 백엔드 개발자를 원하지 않는다고 대차게 떨어졌다.) 

 

삼성 코딩테스트 문제 같은 경우 호흡이 굉장히 길기 때문에 이를 나누고 분석하는 과정이 매우 중요하다고 생각한다. 

---

💡 풀이 과정

• 각 명령에서 필수적으로 필요한 것들을 우선적으로 파악하고 전체적 틀을 잡는 것이 좋다고 생각한다. 
• 노드 추가
  • 가장 기본적으로 해당 노드에서 부모-자식, 색깔, 최대 깊이를 알고 있어야한다.
  • 노드를 상하위로 자유롭게 움직이기 위해서 parents, childs 리스트와 딕셔너리를 같이 관리하자.
  • max_depth 를 검증하기 위해서 노드 추가 시 최상위 노드까지 max_depth 와 실제 depth 를 비교 후 추가한다. 
• 점수 조회
  • 점수 조회를 위해서 각 노드에서 하위 노드에 존재하는 색을 set 형태로 관리하자.
• 색깔 변경
  • 노드의 색을 변경하면서 childs 를 따라 하위 노드들의 색을 모두 변경해야한다. 이때 자신과 하위 노드의 set 은 모두 해당 색 하나의 set 으로 동일하다.
  • 상위 노드의 경우 해당 상위 노드의 자식들의 set 을 합쳐 다시 계산해야한다. 
• 색깔 조회
  • 기록해 놓은 색을 바로 조회한다. 
• parents : 노드는 하나의 부모만을 가지고 있기 때문에 리스트로 관리할 수 있다. 오직 부모 노드 번호만 관리한다.
• childs : { node_id : [ [ childs_id_list ] , max_depth, color, set( colors ) ]
  • 노드의 아이디를 통해 구별되며 자식들의 리스트, 최대 깊이, 색깔, 하위 노드들의 색의 Set 을 가져야한다. 
• parents 와 childs 만 위 처럼 관리하면 해결할 수 있다. 

---

📌 포인트

• 노드는 어떤 노드든 max_depth 를 만족하지 않으면 추가될 수 없다.
• 색깔 변경 시 하위 노드와 상위 노드를 나누어 변경해야한다. 
• 명령은 최대 20,000 번 주어지지만 id 값은 1 <= ,,, <= 100,000 이므로 parents 리스트 생성에 주의하자.

---

💻  코드

# https://www.codetree.ai/training-field/frequent-problems/problems/color-tree/description?page=1&pageSize=5

import sys

input = sys.stdin.readline

parents = [-1] * 100001
childs = {-1:[[], 0, 0, {}]}


def validate_add_node(pId):
    path = []
    depth = 1

    while pId != -1:
        if childs[pId][1] <= depth:
            return []
        depth += 1
        path.append(pId)
        pId = parents[pId]

    return path


def add_node(mId, pId, color, max_depth):
    if pId == -1:
        parents[mId] = pId
        childs[pId][0].append(mId)

        parents[mId] = pId
        if mId not in childs:
            childs[mId] = [[], max_depth, color, {color}]
        else:
            childs[pId][0].append(mId)

    else:
        path = validate_add_node(pId)
        if path:
            for i in range(len(path) - 1, -1, -1):
                childs[path[i]][3].add(color)

            parents[mId] = pId
            childs[pId][0].append(mId)
            if mId not in childs:
                childs[mId] = [[], max_depth, color, {color}]
            else:
                childs[mId][0].append(mId)


def change_color(mId, color):
    stack = [mId]
    while stack:
        m = stack.pop()

        if m not in childs:
            continue

        childs[m][2] = color
        childs[m][3] = {color}

        stack.extend(childs[m][0])

    pId = parents[mId]
    while pId != -1:
        new_set = {childs[pId][2]}
        child = childs[pId][0]
        for c in child:
            new_set.update(childs[c][3])
        childs[pId][3] = new_set

        pId = parents[pId]


def find_color(mId):
    return childs[mId][2]


def calculate_value():
    stack = [-1]
    value = 0
    while stack:
        mId = stack.pop()

        if mId not in childs:
            continue

        value += pow(len(childs[mId][3]), 2)
        stack.extend(childs[mId][0])

    return value


def solution():
    Q = int(input())

    for i in range(Q):
        command = list(map(int, input().split()))

        if command[0] == 100:
            m, p, c, md = command[1:]
            add_node(m, p, c, md)
        elif command[0] == 200:
            mId, color = command[1:]
            change_color(mId, color)
        elif command[0] == 300:
            mId = command[1]
            print(find_color(mId))
        elif command[0] == 400:
            print(calculate_value())

solution()