BOJ 18921번: Cost of Subtree

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

18921번: Cost Of Subtree

일단 영어로 된 문제는 백스탭하고 보지만.. 일단 읽어보았다..

처음에 A subtree of the tree is defined as a non-empty connected subset of its edges. 이 부분이 이해가 안돼서 뭐래는 거야;; 하다가 그냥 말 그대로 연결된 "간선"의 부분집합이구나 깨달았다.

크루스칼 알고리즘을 생각하면 된다는 힌트를 주워듣고(?) 보니 가중치 값이 큰 간선부터 문제가 요구하는 대로 그 간선을 연결하고 그 간선이 속하는 subtree의 코스트를 구해나가면서 정답을 갱신하면 되었다.

간선이 속하는 subtree가 어디인지 구하기 위해 find와 union 함수를 이용하였고, cost라는 배열을 따로 만들어 각 부분집합 별로 간선의 개수와, 제일 작은 가중치 값을 저장했다.(find(x)값을 인덱스로)

그리고, union함수에서 find의 값을 리턴시켜서 연결된 간선의 개수를 더할 때 사용했다.

 

import sys
input = sys.stdin.readline
N = int(input())
edge = []
cost = [[0,float('INF')] for _ in range(N+1)]

for _ in range(N-1):
    a, b, v = map(int, input().split())
    edge.append([a,b,v])

edge.sort(key = lambda x: x[2], reverse = True)
ans = 0
parent = [i for i in range(0,N+1)]

def find(x):
    if x==parent[x]:
        return x
    return find(parent[x])

def union(x,y):
    x = find(x)
    y = find(y)
    if (x<y):
        parent[y] = x
        return x, y
    else:
        parent[x] = y
        return y, x
    
for i in edge:
    #if (find(i[0]) != find(i[1])):
        a, b = union(i[0], i[1])
        cost[a][0] += 1 + cost[b][0]
        cost[a][1] = min(cost[a][1], i[2])
        ans = max(ans, cost[a][0]*cost[a][1])
print(ans)

처음에는 부분집합 구하는 문제처럼 if (find(x) != find(y)): 조건을 써 놓았는데, 생각해보니 이건 트리니까 순환하는 그래프가 나올 수가 없기 때문에 불필요하다 판단되어 뺐다.