본문 바로가기

Python/강좌

Kruskal Algorithm 시뮬레이터

MST를 구하는 유명한 알고리즘인 Kruskal 알고리즘의 시뮬레이터를 만들어 봤다.

 

이론상 노드가 얼마가되든 상관 없지만, 보기가 힘들니 5개 정도만 쓰자.

점이 좀 잘 퍼졌으면 좋겠는데, 그렇지 않아서 아쉽다.

 

프로그램이 끝나면 MST를 출력하거나, 못만들었다는 얘기를 해준다.

이번에도 코드가 구데기니 사용만 하자.

 

간선의 중간쯤에 있는 것은 가중치다.

초록색은 다음에 볼 간선이다.

초록색이 떴고, 빨간색으로 변하면 MST의 간선이 된 것이고, 검은색으로 변하면 사이클이 발생해 MST에 넣지 못한 것이다.

간선 생길확률은, 확률 결정하는 값을 $a$라고 할때, $\frac{3}{a}$다.

그리고 간선 개수는 확률적으로  노드 개수를 $n$이라고 할때, 간선 만드는 루프를 $n-1 + n-2 + \dots + 1$번 반복하고, 이는 등차수열 합 공식에 의해 $\frac{n^2}{2}$와 같다.

따라서, 확률은$\frac{3}{a} \times \frac{n^2}{2} = \frac{3n^{2}}{2a}$다.

import pygame
import random

def find(n):
    if p[n] == n:  return n
    p[n] = find(p[n])
    return p[n]

def union(n1,n2):
    global p
    n1 = find(n1)
    n2 = find(n2)
    if n1 < n2:
        p[n2] = n1
    else:
        p[n1] = n2

def isInSameSet(n1,n2):
    return find(n1) == find(n2)

def getDistSquared(p1,p2):
    return (p1[0] - p2[0])**2 + (p1[1] - p2[1])**2

def dotToScreenDot(a):
    return a[0], SCREEN_SIZE[1]-a[1]

def drawNode(nodeN, color):
    pygame.draw.circle(screen, color, dotToScreenDot(nodes[nodeN]), NODE_RADIUS, 3)
    screen.blit(font.render(f'#{nodeN}', True, color), dotToScreenDot((nodes[nodeN][0] - 5, nodes[nodeN][1] + 5)))

def drawEdge(edgeNum, color):
    v,e,w = edges[edgeNum]
    pygame.draw.line(screen, color, dotToScreenDot(nodes[v]),dotToScreenDot(nodes[e]), 3)
    midPoint = (nodes[v][0] + nodes[e][0])//2, (nodes[v][1] + nodes[e][1])//2
    screen.blit(font.render(str(w), True, BLUE), dotToScreenDot(midPoint))

RED = (255, 0, 0)
BLACK = (0, 0, 0)
BLUE = (0, 0, 255)
WHITE = (255, 255, 255)
GREEN = (0, 255, 0)
SCREEN_SIZE = [1920, 1080]
nodeCount = int(input("노드 개수를 입력하세요>>>"))
FPS = int(input("FPS >>>"))
spaceMode = int(input("스페이스로 다음 단계를 진행하시겠습니까? 스페이스 사용하려면 1, 아니면 0>>>"))
b = int(input("간선 생길 확률 결정하는 것. 1보다 커야하고, 낮을수록 확률이 높아짐>>>"))
pygame.init()
screen = pygame.display.set_mode(SCREEN_SIZE, pygame.FULLSCREEN)
done = False

nodes = [[random.randint(0, SCREEN_SIZE[0]), random.randint(0, SCREEN_SIZE[1])] for _ in range(nodeCount)]
p = [i for i in range(nodeCount)]
edges = []
eA = edges.append
for i in range(nodeCount):
    for j in range(i+1, nodeCount):
        if random.randint(0,b) < 3:
            eA([i, j, random.randint(0, 1000)])
edges.sort(key=lambda x: x[2])
NODE_RADIUS = 20
clickedNode = None
font = pygame.font.Font(None, 30)
clock = pygame.time.Clock()
MST = []
if spaceMode:
    goNext = False
else:
    goNext = True
i = 0
while len(MST) < nodeCount-1:
    screen.fill(WHITE)
    for j in range(nodeCount):
        drawNode(j, BLACK)
    for j in range(len(edges)):
        drawEdge(j, BLACK)
    for j in range(len(MST)):
        drawEdge(j, RED)
    drawEdge(i, GREEN)#다음 선택할 것
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
                pygame.quit()
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_ESCAPE:
                pygame.quit()
            if event.key == pygame.K_SPACE and spaceMode:
                goNext = True
    if goNext:
        a,b,w = edges[i]
        if not isInSameSet(a,b):
            MST.append(i)
            union(a,b)
        i += 1
        if spaceMode:
            goNext = False
    pygame.display.flip()
    clock.tick(FPS)
if len(MST) != nodeCount-1:
    print("Can't Make MST")
else:
    print([edges[i] for i in MST])
pygame.quit()

아래는 사용예시다.

노드 개수는 7개, FPS는 60, 스페이스바로 진행은 1, 간선생길 확률 결정하는 것은 6을 넣었다.

클릭해서 보자.

'Python > 강좌' 카테고리의 다른 글

K-means 시뮬레이터  (0) 2020.09.15
Union-Find 시뮬레이터  (0) 2020.09.14
최소크기 외접원-경사하강법 시뮬레이터  (0) 2020.09.13
Monotone Chain 시뮬레이터  (0) 2020.09.13