[코드트리/삼성SW역량테스트] 코드트리 메신저 (Python)

Problem : https://www.codetree.ai/training-field/frequent-problems/problems/codetree-messenge

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

 

Difficulty : Platinum 4

 

Status : Solved

 

Time : 01:58:36

 

---

 

문제 설명 / 입력 및 출력

 

 

자세한 설명은 코드트리 사이트 링크를 참조해 주세요!

 

 

---

 

풀이

 

구현 + 트리 문제. 구현 자체의 난도는 쉬울 것으로 판단된다. 다만 고려해야 할 쿼리가 딱 하나 존재한다.

 

500 쿼리(본 노드까지 알람이 오는 모든 노드수 탐색)를 그래프 탐색으로 구현한다면, 최악의 경우 O(N)의 시간복잡도가 나타난다. 따라서 최대 시간복잡도는 O(QN)으로 시간 초과를 나타나게 된다.

 

이를 해결하기 위해, 세그먼트 트리와 비슷한 방법을 사용하자. 본 노드에는 현재 노드 시점에 영향을 미치는 알람들과 그 거리들을 저장토록 한다. 최대 깊이가 20이므로, 길이 21의 리스트를 사용하는 것을 추천한다(딕셔너리 등의 구조 역시 동일한 역할을 수행하지만, 생성 및 업데이트에 걸리는 시간을 고려하면 그렇게 좋은 방법은 아니다).

 

알람 설정 변경, 현재 알람 세기 변경, 위치 변경 등을 수행하게 되면, 이 노드뿐 아니라 이 노드의 모든 상위 노드에도 영향을 미치게 된다. 따라서 이런 노드의 변경시마다 본 노드, 혹은 상위 노드부터 재귀적으로 값을 업데이트해 줄 필요가 있다. 어떤 노드를 업데이트할 때는 자식 노드(최대 2개)들의 알람 거리 리스트를 참조할 수 있으며, 자식 노드가 보유하고 있는 거리 i ( i > 0)의 값 j는 부모 노드의 거리 i-1에 반영된다. 이 경우, 이 쿼리를 수행하는 업데이트 작업은 O(QlogN)이 소요된다.

 

 

풀이 코드

from collections import defaultdict
import sys
input = sys.stdin.readline

N, Q = map(int, input().split())
parents = [0]*(N+1)
children = [list() for _ in range(N+1)]
auth = [0]*(N+1)
surplus_auth = [[0]*21 for _ in range(N+1)]
alarms = [1]*(N+1)

def init() :
    global parents, auth
    q = list(map(int, input().split()))
    for i in range(N) :
        parents[i+1] = q[i+1]
        auth[i+1] = min(20, q[i+1+N])
    for i in range(1, N+1) :
        children[parents[i]].append(i)
        power, idx = auth[i], i
        while power > -1 :
            surplus_auth[idx][power] += 1
            power -= 1
            if idx == parents[idx] :
                break
            idx = parents[idx]

def update(idx) :
    cnt = 0
    surplus_auth[idx] = [0]*21
    surplus_auth[idx][auth[idx]] += 1
    for child in children[idx] :
        if not alarms[child] :
            continue
        for i, val in enumerate(surplus_auth[child][1:]) :
            surplus_auth[idx][i] += val
    if idx != parents[idx] :
        update(parents[idx])

def change_alam(c) :
    alarms[c] = 1 - alarms[c]
    update(parents[c])
    
def change_power(c, power) :
    auth[c] = min(20, power)
    update(c)
    
def change_parents(c1, c2) :
    p1, p2 = parents[c1], parents[c2]
    if p1 == p2 :
        return

    parents[c1], parents[c2] = parents[c2], parents[c1]
    children[p1].remove(c1)
    children[p1].append(c2)
    children[p2].remove(c2)
    children[p2].append(c1)

    update(p1)
    update(p2)

def check(c) :
    print(sum(surplus_auth[c]) - 1)

init()
for _ in range(Q-1) :
    query, *body = map(int, input().split())
    if query == 200 :
        c = body[0]
        change_alam(c)
    elif query == 300 :
        c, power = body
        change_power(c, power)
    elif query == 400 :
        c1, c2 = body
        change_parents(c1, c2)
    else :
        c = body[0]
        check(c)

풀이 완료!