(백준-1058) 친구 - 파이썬

문제

백준-1058 친구 문제 보러가기(출처)

문제를 혼자 풀어보려 했으나 문제 이해가 너무 어려워서 다른분의 글을 참조하여 코드를 짜보았습니다. 참고한 사이트는 아래에 적어놓겠습니다!
참고 사이트 : https://jaimemin.tistory.com/615

문제 이해

문제에서 구할 것은 가장 유명한 사람의 2-친구 수를 구하는 것입니다. 2-친구를 이해하는데 좀 시간이 걸렸다는... 이 문제에서 2-친구는 A라는 사람이 B의 2-친구가 되기 위해서는 'A와 B가 친구'이거나 'A와 친구이고 B와 친구인 C가 존재' 이 두가지 경우에 2-친구가 되기 위한 조건을 만족합니다.

결론: 2-친구는 A와 B서로 친구이거나 A와 B가 동시에 친구인 C가 존재할 경우에 만족한다는 것입니다.

입력예시:
5
NYNNN
YNYNN
NYNYN
NNYNY
NNNYN
아래는 위의 입력을 받았을 때 친구관계를 그래프형태로 도식화한 것입니다.

노드별 2-친구관계 수
노드 1: [1 - 2] --> 1개
노드 2: [2 - 1], [2 - 3], [2 - 3 - 4] --> 3개
노드 3: [3 - 2], [3 - 2 - 1], [3 - 4], [3 - 4 - 5] --> 4개
노드 4: [4 - 5], [3 - 4], [1 - 3 - 4] --> 3개
노드 5: [4 - 5] --> 1개

결과적으로 이 친구관계에서 가장 유명한 사람은 3번이고 2-친구의 수는 4명입니다.

접근

그렇다면 우리는 어떻게 문제를 풀어볼수 있을까요?
바로 플로이드 워셜 알고리즘입니다! 그 이유는 입력으로 받는 친구관계를 그래프로 나타낼 수 있기 때문입니다.
들어오는 입력을 살펴보면
5
NYNNN
YNYNN
NYNYN
NNYNY
NNNYN
Y: 친구관계 N: 친구가 아닌 관계

위의 입력을 통해 플로이드 워셜 알고리즘에서 사용하는 최단거리 테이블을 만들 수 있습니다. 친구관계인 경우에는 1로 아닌경우는 INF(무한대), 자기 자신의 경우는 0으로 만든 후 플로이드 워셜 알고리즘을 실행합니다.

코드

import sys
input = sys.stdin.readline

# A가 B와 2-친구가 되기위한 조건
# 1. A와 B가 친구사이이어야 함
# 2. A와 친구이고, B와 친구인 C가 존재해야된다.

# 사람의 수
n = int(input().strip())

# 친구관계 입력받기
friends = []
for _ in range(n):
    friends.append(list(input().strip()))
    
INF = int(1e9)

# 최단거리 테이블 생성(INF는 친구관계가 아닌 경우에 저장된다.)
dist = [[INF] * n for _ in range(n)]

for i in range(n):
    for k in range(n):
        if friends[i][k] == 'Y':
            dist[i][k] = 1
        if i == k:
            dist[i][k] = 0

# 플로이드 워셜 알고리즘을 이용하여 친구관계 사이의 거리를 계산하여 최단거리 테이블을 갱신한다.
def floid():
    for i in range(n):
        for k in range(n):
            for h in range(n):
                if i == k or i == h or k == h:
                    continue
                else:
                    dist[k][h] = min(dist[k][h], dist[k][i] + dist[i][h])

floid()
for i in dist:
    print(i)

# i와 k사이의 거리가 2 이하인 경우는 2-친구의 조건을 만족하기 때문에 갯수를 센다
result = 0
for i in range(n):
    cnt = 0
    for k in range(n):
        if i == k:
            continue
        if dist[i][k] <= 2:
            cnt += 1
    result = max(result, cnt)
print(result)

코드는 위와 같이 짰습니다.
1. 문제에서 입력을 받습니다.
2. 입력받은 값들을 기반으로 최단거리테이블을 초기화합니다.(자기자신: 0, 친구관계: 1, 친구관계X: INF)
3. 초기화한 최단거리테이블을 가지고 플로이드 워셜 알고리즘을 적용하여 최단거리 테이블을 갱신합니다.
4. 최단거리 테이블을 행별로 2를 초과하지 않는 수를 세면서 행을 넘어갈때마다 최댓값을 갱신해줍니다.
5. 4에서 구한 최댓값이 우리가 구할 최댓값입니다!

후기

문제를 이해하기 좀 어려웠고 읽자마자 플로이드 워셜 알고리즘 써야겠다는 생각이 들지 않았습니다... 벽이 느껴진달까? 플로이드 워셜 알고리즘을 공부해야할 필요성을 느끼게 되었다..