[백준] 4315 - 나무 위의 구슬 (Python)

알고리즘 (PS)/BOJ

[백준] 4315 - 나무 위의 구슬 (Python)

에버듀 2024. 11. 12. 20:16

https://www.acmicpc.net/problem/4315

재미있는 트리 연습 문제

아이디어 + 구현 모두 깐깐한 문제였다.

각 노드가 갖고 있는 구슬 개수가 모두 1개가 되도록 할 때 구슬을 옮기는 최소 횟수를 구하는 문제이다.

내가 생각한 풀이 흐름은 다음과 같다.

1. 전체 트리를 DFS 돌면서 현재 자신을 root 로 하는 서브트리에 대해 이 서브 트리가 가지고 있는 모든 구슬의 개수, 필요한 구슬의 개수를 구한다. 만약 가지고 있는 구슬이 필요한 구슬보다 많다면, 그 차이를 return 하여 부모 노드로 여분 구슬을 넘긴다.

DFS 실행이 끝나면 모든 노드는 자신을 루트로 하는 서브트리 모두를 채울 수 있도록 딱 맞게 구슬을 갖고 있거나, 모자라거나 둘 중 하나의 상태가 된다.

2. 두번째로 다시 DFS를 돌면서 각 노드가 가진 여분 구슬을 재분배한다. 필요 구슬 > 가진 구슬인 자식 노드에 대해 그 차이만큼 자신의 여분 구슬을 나누어주면 된다. 이때 중요한 것은 자식 노드가 필요 구슬 합 = 가진 구슬 합 이더라도 분배를 하도록 호출해야 한다.

그 안에서 내부적으로 재분배가 필요할 수 있기 때문이다.

import sys
input = sys.stdin.readline
sys.setrecursionlimit(10**5)

def check_count(node):
    global answer

    now_need_count = 1
    now_have_count = marbles[node]

    for conn_node in graph[node]:
        if not visit[conn_node]:
            visit[conn_node] = True
            check_need_count, check_have_count = check_count(conn_node)
            if check_need_count < check_have_count:
                answer += (check_have_count - check_need_count)
                marbles[node] += (check_have_count - check_need_count)
                marbles[conn_node] = check_need_count

            now_have_count += check_have_count
            now_need_count += check_need_count

    need_count[node] = now_need_count
    have_count[node] = now_have_count

    return now_need_count, now_have_count


def distribute_marbles(node):
    global answer
    for conn_node in graph[node]:
        if not visit[conn_node]:
            visit[conn_node] = True
            if need_count[conn_node] > have_count[conn_node]:
                answer += (need_count[conn_node] - have_count[conn_node])
                marbles[node] -= (need_count[conn_node] - have_count[conn_node])
                marbles[conn_node] += (need_count[conn_node] - have_count[conn_node])
            distribute_marbles(conn_node)


while True:
    n = int(input())
    if n == 0:
        break

    graph = [[] for _ in range(n+1)]
    visit = [False for _ in range(n+1)]
    marbles = [0]*(n+1)
    need_count = [0]*(n+1)
    have_count = [0]*(n+1)
    answer = 0
    for _ in range(n):
        node, marble, *children = map(int, input().split())
        graph[node] += children[1:]
        for conn_node in children[1:]:
            graph[conn_node].append(node)
        marbles[node] += marble

    visit[1] = True
    check_count(1)

    visit = [False for _ in range(n + 1)]
    visit[1] = True
    distribute_marbles(1)

    print(answer)

구현이 진짜 어려웠는데, 풀고나니 다른 풀이로는 필요구슬 - 가진 구슬의 절댓값을 다 더하면 그냥 정답이 된다는 걸 보고 저렇게 하면 더 간단하구나 싶었다.

역시 수학을 잘해야 손이 덜 고생한다..

문제를 풀고서 틀렸길래 테스트 케이스를 찾아봐야겠다 싶어서 문제 출처를 찾아봤더니 워털루 대학교 대회 문제였다.

https://uwaterloo.ca/international-collegiate-programming-contest/past-local-contest-results

Past Local Contest Results | International Collegiate Programming Contest

3 October, 2021-Waterloo local contest Zixiang Zhou won - Final standings. Problem Set

uwaterloo.ca

여기에 모든 대회 문제와 테스트케이스가 존재한다.

그런데 막상 접속해봤더니 DNS 에러가 뜨면서 접속이 되지 않았다.

가장 최근에 열린 2024년 대회도 접속이 안되는 것을 보고 이상하다고 생각했다.

처음엔 워털루 대학교 학생만 접근할 수 있도록 막았나 싶었는데, 보통 테스트 케이스를 공개할 때 그렇게 공개하는 경우는 못봤어서 한번 최근에 진행중인 대회 사이트에 들어가보았다.

https://icpc.student.cs.uwaterloo.ca/icpc/index.php