백준 - 구슬 찾기(2617)

문제 링크 ▶︎ 백준 구슬찾기 2617

문제 파악

무게 순으로 따졌을 때, 절대로 중간에 오지 않는 구슬을 찾는다는 건 나보다 가벼운 구슬이 n/2개 이상있거나, 나보다 무거운 구슬이 n/2개 이상 있는 구슬을 찾는 것과 동일하다고 생각했다.

그래서 처음부터 입력을 받을 때, 나보다 무거운 구슬을 입력받는 해시와 나보다 가벼운 구슬을 입력받는 해시 두개로 나눠서 접근하면 어렵지 않게 풀 수 있을 것 같았다.

가벼운 구슬 찾고, 무거운 구슬을 찾는 식으로 두번 계산하는 것보다 메모리를 더 써서 동시에 찾는 방법으로 접근해보자.

접근 방법

1. more이라는 해시에는 나보다 무거운 구슬을 받고, less라는 해시에는 나보다 가벼운 구슬을 저장해둔다. 또한 more_visit, less_visit 을 각각 만들어둔다.

2. dfs를 찾는 것은 역시 1번부터 시작해서 n번까지로 진행하며, 각각 해시에서 나보다 무거운, 나보다 가벼운 구슬을 찾기 위해서 인자를 넘겨줘가면서 재귀함수를 수행한다. 그리고 방문하게 되면 조건에 맞는 인덱스 visit을 1로 변경해준다.

3. 즉 1번 구슬부터 시작해서 1번 구슬보다 가벼운 구슬을 탐색하면 해당 인덱스의 less_visit이 1로 변경되었을 것이고 무거운 구슬을 탐색하는 과정이 끝나면 해당 인덱스의 more_visit이 1로 변경되었을 것이다.

4. 그렇다면 sum(visit) 이 해당 구슬보다 무거운 구슬의 수, 가벼운 구슬의 수가 될 것이기 때문에 n/2 보다 크다면 answer를 올려주고, 다 돌고나서 answer를 출력한다.

코드

import sys
import math
sys.setrecursionlimit(10**5)
input = sys.stdin.readline

def dfs(balls,num,visit):
    for idx in balls[num]:
        if visit[idx] == 0:
            visit[idx] = 1
            dfs(balls,idx,visit)

n, m = map(int,input().split())
more = {i:[] for i in range(n+1)}
less = {i:[] for i in range(n+1)}
answer = 0
for _ in range(m):
    a, b = map(int,input().split())
    less[a] += [b]
    more[b] += [a]
for i in range(1,n+1):
    more_visit = [0] * (n+1)
    less_visit = [0] * (n+1)
    dfs(more,i,more_visit)
    dfs(less,i,less_visit)
    if sum(more_visit) > (n//2) or sum(less_visit) > (n//2):
        answer += 1

print(answer)

개선 사항

잘 짠 코드는 아닌 것 같다. 해시도 두개를 사용했고, 방문 배열도 두개를 사용했기 때문에 더 좋은 아이디어를 충분히 고민해볼 수 있다고 생각한다.