[Python/파이썬] 백준 2458번 키 순서

2458번: 키 순서

 

코드

처음 풀이

import sys
input = sys.stdin.readline

n, m = map(int, input().split())
dist = [[[int(1e9), -int(1e9)] for _ in range(n+1)] for _ in range(n+1)]
for i in range(n+1):
    dist[i][i] = [0, 0]
for _ in range(m):
    a, b = map(int, input().split())    # a < b
    dist[a][b][0] = 1
    dist[b][a][1] = -1

for k in range(1, n+1):
    for i in range(1, n+1):
        for j in range(1, n+1):
            if dist[i][j][0] > dist[i][k][0]+dist[k][j][0]:
                dist[i][j][0] = dist[i][k][0]+dist[k][j][0]
            if dist[i][j][1] < dist[i][k][1]+dist[k][j][1]:
                dist[i][j][1] = dist[i][k][1]+dist[k][j][1]

ans = n
for i in range(1, n+1):
    for j in range(1, n+1):
        if i == j:
            continue
        if dist[i][j][0] == int(1e9) and dist[i][j][1] == -int(1e9):
            ans -= 1
            break
print(ans)

 

처음에는 위와 같이 풀이했다.

 

dist[i][j][0] : i<j인 경우에 i에서 j까지의 거리

dist[i][j][1] : i>j인 경우에 i에서 j까지의 거리

 

이렇게 해두고 처음에 학생들의 키 비교 결과를 입력 받을 때 큰 경우는 1, 작은 경우는 -1을 넣어줬다. a<b인 경우에는 dist[a][b][0] = 1, dist[b][a][1] = -1 이렇게 말이다.

그리고 플로이드 워셜 알고리즘을 이용하여 모든 노드에서 다른 모든 노드까지의 최단 거리를 구한다. 이때 그래프는 방향성이 있는 그래프로 현재보다 키가 큰 노드로 가다가 작은 노드로 가거나 현재보다 작은 노드쪽으로 가다가 큰 노드로 가는 것은 안된다.

위와 같은 예시가 안된다는 말이다. 계속 현재보다 큰 노드로 이동하거나 작은 노드로 이동하는 것만 된다. 이렇게 구하면 i->j로 갈 때 큰 노드로만 이동해서도 갈 수 없고, 작은 노드로만 이동해서도 갈 수 없는 경우가 생기는데 이런 경우가 i와 j의 키를 비교할 수 없는 경우이다. 그러므로 플로이드 워셜을 수행한 뒤에 dist[i][j]가 [int(1e9), -int(1e9)]가 하나라도 있는 i행은 정확한 순서를 알 수 없는 노드이다.

 

근데 이렇게 하면 너무 복잡한 것 같아서 고쳤다.

 

수정한 풀이

import sys
input = sys.stdin.readline

n, m = map(int, input().split())
graph = [[0]*(n+1) for _ in range(n+1)]
for _ in range(m):
    a, b = map(int, input().split())    # a < b
    graph[a][b] = 1

for k in range(1, n+1):
    for i in range(1, n+1):
        for j in range(1, n+1):
            if graph[i][k] == 1 and graph[k][j] == 1:
                graph[i][j] = 1  # i < j

ans = 0
for i in range(1, n+1):
    known = 0
    for j in range(1, n+1):
        known += graph[i][j]+graph[j][i]
    if known == n-1:
        ans += 1
print(ans)

 

graph[i][j]는 i < j인 경우에 1이고, 아니거나 모르는 경우에는 0이다.

graph[i][j]+graph[j][i] = 1이라면 i와 j는 서로의 비교 결과를 알 수 있는 경우이다. 그러므로 노드 i를 기준으로 잡고 다른 모든 노드에 대해 graph[i][j]+graph[j][i]의 값의 합을 구했을 때 n-1이라면 i는 다른 모든 노드들과의 비교 결과를 알 수 있는 것이므로 정확한 순위를 알 수 있다.

 

이렇게 풀이하고 제출한 결과, 처음의 풀이보다 시간을 많이 단축시킬 수 있었다.

 

 

근데 두 풀이 다 Python3에서는 시간 초과가 발생하므로 PyPy3로 제출해야 한다. 아마 플로이드 워셜의 시간 복잡도가 O(n^3)이기 때문인 것 같다.