백준 14621 : 나만 안되는 연애 (Python)

본문 링크

📌 크루스칼 알고리즘

import sys
input=sys.stdin.readline

"""
남초대학교-남초대학교 or 여초대학교-여초대학교로 연결되면 안됨.
연결안된것은 set 으로 판별.
"""

def Find(x):

    if x!=disjoint[x]:
        disjoint[x]=Find(disjoint[x])
    return disjoint[x]


N,M=map(int,input().split())

disjoint=[ _ for _ in range(N+1) ]

School=list(map(str,input().split()))

edges=[] ; total=0

for i in range(M):

    u,v,d=map(int,input().split())
    edges.append((u,v,d))

edges.sort(key=lambda x:x[2]) # 간선의 가중치에 대해 정렬

for x,y,value in edges:

    if School[x-1]!=School[y-1]:
        x=Find(x)
        y=Find(y)
        if x!=y:
            if x>y:
                disjoint[x]=y
            else:
                disjoint[y]=x
            total+=value

check=set()

for i in range(1,N+1):
    if Find(i) not in check:
        check.add(Find(i))

if len(check)>1:
    print(-1)
else:
    print(total)

📌 프림 알고리즘

import sys,heapq
input=sys.stdin.readline

N,M=map(int,input().split())

School=list(map(str,input().split()))

Tree=[ [] for _ in range(N+1) ]

for i in range(M):

    u,v,d=map(int,input().split())

    Tree[u].append((d,v,School[v-1]))
    Tree[v].append((d,u,School[u-1])) # 가중치 먼저


visit=[False]*(N+1) ; heap=[(0,1,School[0])] ; total=0
while heap:

    value,node,color=heapq.heappop(heap)

    if not visit[node]:
        visit[node]=True
        total+=value
        for i,j,tree_color in Tree[node]:
            if color!=tree_color:
                heapq.heappush(heap,(i,j,tree_color))


if visit[1:].count(False)>=1:
    print(-1)
else:
    print(total)

📌 어떻게 접근할 것인가?

일반적인 최소 스패닝 트리이지만 한가지 조건이 추가되었습니다. 같은 색깔(대학) 은 서로 연결 할 수 없습니다.

따라서 크루스칼 알고리즘에서는 if School[x-1]!=School[y-1] 라는 조건을 한번 체크 해주고 간선을 연결하고 연결할수 있는지 없는지에 대한 여부는

check=set()

for i in range(1,N+1):
    if Find(i) not in check:
        check.add(Find(i))

if len(check)>1:
    print(-1)
else:
    print(total)

위 코드를 통해 확인했습니다.

프림 알고리즘에서는 간선을 입력받을때 대학의 색깔도 같이 입력받고 if color!=tree_color 라는 조건을 추가하여서 $Tree$ 에서 원소를 추가할때 서로 색깔이 다른 정점들만 heap 에 원소를 추가해주었습니다.

프림 알고리즘에서 모든 정점이 연결되어 있는가는 visit 배열을 보면되는데

if visit[1:].count(False)>=1:
    print(-1)
else:
    print(total)

위와 같은 코드로 확인 가능합니다.