HackerRank / Kruskal (MST) : Really Special Subtree / Python

 

https://www.hackerrank.com/challenges/kruskalmstrsub/problem?isFullScreen=true

Kruskal (MST): Really Special Subtree | HackerRank

참조

https://techblog-history-younghunjo1.tistory.com/262

 

 

크루스칼 알고리즘에 들어가기 앞서 신장 트리, 최소 신장 트리에 대해 알아야 한다.

위의 블로그에 잘 정리가 되어있지만 짧게 설명을 하자면

1) 신장 트리 - 하나의 그래프가 있을 때, 모든 노드를 포함하면서 즉, 모든 노드들 간에 서로 연결은 되어있지만 사이클이 존재하지 않는 '부분' 그래프이다.

2) 최소 신장 트리 - 간선 비용이 존재할 때 사이클이 형성되지 않으며 간선 비용이 최소가 되는 알고리즘이다.

 

문제 특이점

위의 문제의 경우 크루스칼 알고리즘이나 간선 비용이 같을시 '출발 노드 + 도착 노드 + 간선 비용'이 작은 edge를 사용하라고 나와있다.

 

문제 접근

1) weight에 대해 edge 정렬

2) 정렬된 리스트에 대해 sum값으로 부분 정렬

3) Union_Find를 통한 문제 해결

 

#!/bin/python3

import math
import os
import random
import re
import sys

def kruskals(g_nodes, g_from, g_to, g_weight):
    sort_list = [] # 정렬할 리스트 생성
    length = 0
    for i in range(len(g_from)): # weight에 대해 정렬
        g_sum = g_from[i] + g_to[i] + g_weight[i]
        sort_list.append([g_from[i],g_to[i],g_weight[i],g_sum])
    sort_list.sort(key = lambda x:x[2])
    
    for s in range(len(sort_list)-1): # 정렬된 리스트에서 sum값에 대해 부분 정렬
        s1 = sort_list[s]
        s2 = sort_list[s+1]
        if (s1[2] == s2[2] & s1[3]>s2[3]):
            tmp = sort_list[s+1]
            sort_list[s+1] = sort_list[s]
            sort_list[s] = tmp
        else:
            pass        
    
    parent = [0 for i in range(g_nodes + 1)] #원래는 -1이나 해당 문제는 노드가 1부터 시작하기에 문제 없다
    for i in range(1, g_nodes+1):
        parent[i] = i
    def find_parent(parent, x): # 부모 노드 확인
        if parent[x] != x:
            parent[x] = find_parent(parent, parent[x])
        return parent[x]

    def union_parent(parent, a, b): # 부모 노드 변경
        a = find_parent(parent, a)
        b = find_parent(parent, b)
        if a < b:
            parent[b] = a
        else:
            parent[a] = b

    for i in range(len(sort_list)): # 부모 노드가 다를 때(=사이클이 생성되지 않을때) union 연산 실행
        if find_parent(parent, sort_list[i][0]) != find_parent(parent, sort_list[i][1]):
            union_parent(parent, sort_list[i][0], sort_list[i][1])
            length+=sort_list[i][2]
    
    return length

if __name__ == '__main__':
    fptr = open(os.environ['OUTPUT_PATH'], 'w')

    g_nodes, g_edges = map(int, input().rstrip().split())
    # print(g_edges)
    g_from = [0] * g_edges
    g_to = [0] * g_edges
    g_weight = [0] * g_edges

    for i in range(g_edges):
        g_from[i], g_to[i], g_weight[i] = map(int, input().rstrip().split())
    
    res = kruskals(g_nodes, g_from, g_to, g_weight)
    
    # Write your code here.
    fptr.write(str(res)+'\n')
    
    fptr.close()