Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

파이썬으로 불행한 친구의 수를 세는 프로그램 구현하기

문제 개요

짝수 명수인 n명의 서로 다른 친구에 대한 선호도 목록이 주어졌다고 가정해 보겠습니다. 각 사람 i에 대해 preferences[i]는 선호 순서대로 정렬된 친구 목록을 담고 있으며, 목록에서 앞쪽에 위치한 친구일수록 더 선호되는 친구입니다. 각 목록의 친구들은 0부터 n-1까지의 정수로 번호가 매겨집니다.

모든 친구는 서로 다른 쌍으로 나뉘며, pairs[i] = [x, y]는 x와 y가 서로 짝을 이루었음을 의미합니다. 이때 어떤 친구 x는 다음 두 조건을 모두 만족하는 친구 u가 존재할 경우 '불행한(unhappy)' 상태가 됩니다.

  • x는 자신의 파트너 y보다 u를 더 선호하고,
  • u 역시 자신의 파트너 v보다 x를 더 선호하는 경우

우리의 목표는 이런 불행한 친구가 총 몇 명인지 구하는 것입니다.

예를 들어, preferences = [[1, 2, 3], [3, 2, 0], [3, 1, 0], [1, 2, 0]], pairs = [[0, 1], [2, 3]]이 주어지면 결과는 2가 됩니다. 0번 친구는 1번과 짝을 이루었지만 3번을 더 선호하고, 3번 역시 2번보다 1번을 더 선호하기 때문에 불행합니다. 마찬가지로 3번 친구는 2번과 짝을 이루었지만 1번을 더 선호하고, 1번도 0번보다 3번을 더 선호하므로 불행합니다.

해결 접근 방법

이 문제는 인접 리스트 형태의 그래프를 활용하면 효율적으로 해결할 수 있습니다. 핵심 아이디어는 각 사람이 자신의 파트너보다 더 선호하는 사람들을 미리 그래프에 기록해 두는 것입니다. 구체적인 단계는 다음과 같습니다.

  1. 그래프를 저장할 빈 인접 리스트(딕셔너리)를 준비합니다.
  2. pairs의 각 쌍 (s, e)에 대해 다음을 수행합니다.
    • preferences[s]를 순회하다가 파트너 e를 만나면 반복을 중단하고, 그 전까지 등장한 선호 대상들을 graph[s]에 기록합니다.
    • preferences[e]에 대해서도 같은 방식으로 처리하여 graph[e]에 기록합니다.
  3. 불행한 친구를 세기 위한 변수 unhappy를 0으로 초기화합니다.
  4. 다시 각 쌍 (s, e)에 대해 다음을 검사합니다.
    • graph[s]에 기록된 선호 대상 pref 중 graph[pref]에 s가 존재하면 s는 불행한 것이므로 unhappy를 1 증가시키고 반복을 중단합니다.
    • e에 대해서도 동일한 방식으로 검사합니다.
  5. 최종 unhappy 값을 반환합니다.

그래프 구축 단계에서 각 쌍마다 선호 목록을 한 번씩만 순회하므로, 전체 시간 복잡도는 O(n²)이며 공간 복잡도 또한 O(n²)입니다.

파이썬 구현 예제

아래 구현 예제를 통해 동작 과정을 더 잘 이해할 수 있습니다.

from collections import defaultdict
def solve(preferences, pairs):
    graph = defaultdict(dict)
    for start, end in pairs:
        for pref in preferences[start]:
            if pref == end:
                break
            graph[start][pref] = 1
        for pref in preferences[end]:
            if pref == start:
                break
            graph[end][pref] = 1

    unhappy = 0

    for start, end in pairs:
        for pref in graph[start]:
            if graph[pref].get(start, None):
                unhappy += 1
                break
        for pref in graph[end]:
            if graph[pref].get(end, None):
                unhappy += 1
                break
    return unhappy

preferences = [[1, 2, 3], [3, 2, 0], [3, 1, 0], [1, 2, 0]]
pairs = [[0, 1], [2, 3]]
print(solve(preferences, pairs))

입력

[[1, 2, 3], [3, 2, 0], [3, 1, 0], [1, 2, 0]], [[0, 1], [2, 3]]

출력

2