[462] 이분 매칭 알고리즘(파이썬 코드)

위와 같은 이분 그래프가 있을 때 이분 매칭을 구해봅시다.

 

1. 종만북 참고 소스코드

종만북 코드는 C언어이므로 파이썬 코드로 옮기면 다음과 같습니다.

n, m = 5, 5
adj = [
    [1, 1, 1, 1, 1],
    [1, 0, 0, 0, 0],
    [1, 0, 0, 0, 0],
    [1, 0, 0, 0, 0],
    [1, 0, 0, 0, 1],
    ]
a_match, b_match = [-1, -1, -1, -1, -1], [-1, -1, -1, -1, -1]
def dfs(a):
    if visited[a]: return False
    visited[a] = True
    for b in range(m):
        if adj[a][b]:
            if b_match[b] == -1 or dfs(b_match[b]):
                a_match[a] = b
                b_match[b] = a
                return True
    return False
size = 0
for s in range(n):
    visited = [False, False, False, False, False]
    if dfs(s): 
        size += 1
        
print(size)
print(a_match, b_match)
# 3
# [1, 0, -1, -1, 4] [1, 0, -1, -1, 4]

하지만 위 코드로 백준에서 이분 매칭 문제를 풀면 시간 초과가 나옵니다.

통과한 사람들의 코드를 참고하여 코드를 살짝 고치면 다음과 같습니다.

 

2. 개선된 이분 매칭 소스코드

n, m = 5, 5
adj = [
    [1, 1, 1, 1, 1],
    [1, 0, 0, 0, 0],
    [1, 0, 0, 0, 0],
    [1, 0, 0, 0, 0],
    [1, 0, 0, 0, 1],
    ]
b_match = [-1, -1, -1, -1, -1]
def dfs(a):
    if visited[a]: return False
    visited[a] = True
    for b in range(m):
        if adj[a][b]:
            if b_match[b] == -1:
                b_match[b] = a
                return True
    for b in range(m):
        if adj[a][b]:
            if not visited[b_match[b]] and dfs(b_match[b]):
                b_match[b] = a
                return True
    return False
size = 0
for s in range(n):
    visited = [False, False, False, False, False]
    if dfs(s): 
        size += 1
        
print(size)
print(b_match)
# 3
# [1, 0, -1, -1, 4]

개선된 부분을 보자면, b에 매칭이 된 정점을 만났을 때 바로 dfs를 도는 것이 아니라 일단 매칭되지 않은 b를 모두 찾은 뒤 없으면 dfs를 돕니다.

dfs를 돌 때도 바로 재귀 함수로 들어가는 것이 아니라, 방문한 적이 있다면 굳이 재귀 함수로 들어가지 않고 다음 정점으로 건너갑니다.

마지막으로 이분 매칭 값을 구할 때 a에서 매칭된 b값은 필요가 없으므로 a_match는 삭제 가능합니다.(만약 a에서 매칭된 b값이 필요하면 넣으면 됩니다.)