[프로그래머스] 중력작용

[Lv.5] 월간 코드 챌린지 시즌2 : 중력작용

문제설명

1부터 n까지 번호가 하나씩 붙은 n개의 노드를 갖는 트리가 주어집니다. 각 노드에는 값이 하나씩 들어 있으며, 이 트리의 루트 노드는 1번 노드입니다. 당신은 이 트리에 대해 다음과 같은 쿼리 두 종류를 처리하면 됩니다.

• 1번 쿼리: 정수 u가 주어집니다. u번 노드의 서브 트리의 모든 노드의 값의 합을 구해야 합니다.
• 2번 쿼리: 정수 u, w가 주어집니다. u번 노드의 값을 삭제한 뒤, u번 노드의 부모 노드의 값을 u번 노드로 복사합니다., u번 노드의 부모 노드에 대해 같은 작업을 반복하며 루트노드까지 거슬러 올라갑니다. 마지막으로 루트 노드의 값을 w로 바꿉니다.

트리의 노드 초기값이 담긴 정수 배열 values, 트리의 연결 상태가 담긴 2차원 정수 배열 edges, 쿼리들이 담긴 2차원 정수 배열 queries가 주어집니다. 쿼리들을 순서대로 처리할 때, 각 1번 쿼리에 대한 답을 수행 순서대로 배열에 담아 return 하도록 solution 함수를 완성해주세요.

제한사항

입출력 예

입출력 예 설명

입출력 예 #1

• 주어진 예시는 1번 쿼리 15개와 2번 쿼리 2개로 이루어져 있습니다.
• 다음 그림은 두개의 2번 쿼리에 의해 트리의 노드 값들이 바뀌는 과정을 나타낸 것입니다.

나의 풀이

접근

• 언어 : Python3
• 접근 : queries 로 쿼리1번을 수행해야하는지 2번을 수행해야하는지 먼저 확인해준다. 쿼리 1번의 경우 서브 트리의 모든 노드의 값의 합을 구해야하므로, 현재 노드를 add_node로 지정하고, 이 노드의 하위 노드를 찾기 위해 edges를 탐색하고, 그 노드의 값을 values에서 찾아 add_result에 더해준다. 마지막으로, add_node 자체의 값을 더해준다.
  쿼리 2번의 경우 부모노드의 값을 내려받고,루트 노드의 값을 주어진 값으로 넣어주는것이다. 주어진 노드 번호를 u로 저장하고, edges에서 u가 하위노드로 지정되어있는 그 부모노드를 찾아 그 값을 바꾸어준다. 그리고 다시 u를 설정해주어 반복한다. 마지막으로 루트노드의 값을 queries에서 입력받은 값으로 넣어준다.

코드

def solution(values, edges, queries):
    answer = []
    add_node = 0
    add_result = 0
    u = 0
    for i in range(len(queries)):
        if queries[i][1] == -1:
            add_node = queries[i][0]
            for j in range(len(edges)):
                if edges[j][0] == add_node :
                    add_result += values[edges[j][1]-1]
            add_result += values[add_node - 1]
            answer.append(add_result)
        else:
            u = queries[i][0]
            for k in range(len(edges)):
                if edges[k][1] == u :
                    values[u - 1] = values[edges[k][0]-1]
                    u = edges[k][0]
            values[u - 1] = queries[i][1]
                     
    return answer

결과

오류 분석

위의 결과 값에서 입력값으로 시뮬레이션을 돌려보면 내가 짠 코드의 오류를 분석할 수 있다.
첫번째 결과값이 11111이 아닌 111이 나온 이유는 쿼리 1번을 짠 코드에서 찾을 수 있는데, 하위노드라는것을 한 단계의 하위노드만 탐색했기 때문이다. 모든 하위노드를 찾기 위해서는 현재 노드값을 계속 바꾸어주어 루프를 돌려야 함을 깨달았다.