SWEA 1251 하나로 (파이썬)

https://swexpertacademy.com/main/code/problem/problemDetail.do?contestProbId=AV15StKqAQkCFAYD 

SW Expert Academy

 

---

풀이

import math

def find_set(x):
    while x != p[x]:
        x = p[x]
    return x

for tc in range(int(input())):
    N = int(input())
    info = [[0, 0] for _ in range(N+1)]
    r_info = list(map(int, input().split()))
    c_info = list(map(int, input().split()))
    for _ in range(N):
        dis_r = r_info[_]
        dis_c = c_info[_]
        info[_+1][0] = dis_r
        info[_+1][1] = dis_c
    E = float(input())
    p = [i for i in range(N+1)]
    edges = []
    for fir in range(1, N+1):
        for sec in range(fir+1, N+1):
            u = fir
            v = sec
            cost = 0
            if info[fir][0] == info[sec][0]:
                cost = abs(info[fir][1] - info[sec][1])
            elif info[fir][1] == info[sec][1]:
                cost = abs(info[fir][0] - info[sec][0])
            else:
                cost = math.sqrt((abs(info[fir][0] - info[sec][0]))**2 + (abs(info[fir][1] - info[sec][1]))**2)
            edges.append([u, v, cost])

    edges.sort(key=lambda x:x[2], reverse=True)

    # for line in edges:
    #     print(line, line[2]**2)

    mst = []
    rst = 0
    cnt = N-1
    while cnt:
        u, v, weight = edges.pop()
        a = find_set(u)
        b = find_set(v)
        if a == b: continue
        rst += (weight**2)*E
        mst.append([u, v, weight])

        p[a] = b
        cnt -= 1

    print(f'#{tc+1}', round(rst))

 

---

후기 및 해설

• 프림이나 크루스칼 알고리즘을 알고 있다면 쉽게 풀 수 있다.
• 나는 크루스칼로 풀었지만, 프림으로 풀면 시간이 더 적게 든다고 배웠다.
• 약간 까다로웠던 부분은 [u, v, weight]를 만들어야 하는데, 구현 문제라 어찌저찌 조건달아서 해놨다.
• 실행시간이 2초대로 나왔는데, 구현에서 시간이 오래걸린건지, 크루스칼이라 오래걸린건지 잘 모르겠다.
• 맘먹고 줄일려면 줄일 수 있겠지만, pass했으니 pass하기로..