본문 바로가기

CodingTest/SW Expert Academy

[ SW Expert Academy ] 5249. 최소 신장 트리

 

1차 시도 [ 8 / 10 ] Fail

헐 .. 시간초과.. 뭐지..

앞에서부터 빼서 그런가..

def make_set(x, p):
    p[x] = x
    return p

def find_set(x, p):
    if p[x] == x:
        return x
    return find_set(p[x], p)

def union(u, v, p):
    if u > v:
        p[find_set(u, p)] = find_set(v, p)
    else:
        p[find_set(v, p)] = find_set(u, p)
    return p

def mst_kruskal(g, n):
    p = {}
    for i in range(n):
        p = make_set(i, p)
    g.sort(key=lambda x:x[2])
    mst_cost = 0
    while g:
        u, v, weight = g.pop(0)
        if find_set(u, p) != find_set(v, p):
            p = union(u, v, p)
            mst_cost += weight

    return mst_cost

T = int(input())
for test_case in range(1, T+1):
    answer = 0
    v, e = map(int, input().split())
    nodes = [tuple(map(int, input().split())) for _ in range(e)] # (n1, n2, w)

    answer = mst_kruskal(nodes, e)

    print("#{} {}".format(test_case, answer))

 

 

2차 시도 [ 10 / 10 ] Pass

와 진짜 그러네...

pop(0)을 pop()으로 바꿔줬고, 이 변화에 맞춰서 정렬을 반대방향으로 설정했다.

g.sort(key=lambda x: x[2])  =>  g.sort(key=lambda x: -x[2])

def make_set(x, p):
    p[x] = x
    return p

def find_set(x, p):
    if p[x] == x:
        return x
    return find_set(p[x], p)

def union(u, v, p):
    if u > v:
        p[find_set(u, p)] = find_set(v, p)
    else:
        p[find_set(v, p)] = find_set(u, p)
    return p

def mst_kruskal(g, n):
    p = {}
    for i in range(n):
        p = make_set(i, p)
    g.sort(key=lambda x:-x[2])
    mst_cost = 0
    while g:
        u, v, weight = g.pop()
        if find_set(u, p) != find_set(v, p):
            p = union(u, v, p)
            mst_cost += weight

    return mst_cost

T = int(input())
for test_case in range(1, T+1):
    answer = 0
    v, e = map(int, input().split())
    nodes = [tuple(map(int, input().split())) for _ in range(e)] # (n1, n2, w)

    answer = mst_kruskal(nodes, e)

    print("#{} {}".format(test_case, answer))

 

 

크루스칼.. 기억 안 나서 공부해서 외워서 적용시켰다 ㅠㅜ