[Python] 백준 16566 - 카드 게임

https://www.acmicpc.net/problem/16566

문제

철수와 민수는 카드 게임을 즐겨한다. 이 카드 게임의 규칙은 다음과 같다.

1. N개의 빨간색 카드가 있다. 각각의 카드는 순서대로 1부터 N까지 번호가 매겨져 있다. 이 중에서 M개의 카드를 고른다.
2. N개의 파란색 카드가 있다. 각각의 카드는 순서대로 1부터 N까지 번호가 매겨져 있다. 이 중에서 빨간색에서 고른 번호와 같은 파란색 카드 M개를 고른다.
3. 철수는 빨간색 카드를 가지고 민수는 파란색 카드를 가진다.
4. 철수와 민수는 고른 카드 중에 1장을 뒤집어진 상태로 낸다. 그리고 카드를 다시 뒤집어서 번호가 큰 사람이 이긴다. 이 동작을 K번 해서 더 많이 이긴 사람이 최종적으로 승리한다. 한 번 낸 카드는 반드시 버려야 한다.

철수는 뛰어난 마술사이기 때문에 본인이 낼 카드를 마음대로 조작할 수 있다. 즉, 카드를 버리고 민수 몰래 다시 들고 온다거나 민수한테 없는 카드를 내기도 한다.

민수는 뛰어난 심리학자이기 때문에 철수가 낼 카드를 알아낼 수 있다. 그래서 민수는 철수가 낼 카드보다 큰 카드가 있다면 그 카드들 중 가장 작은 카드를 내기로 했다.

K번 동안 철수가 낼 카드가 입력으로 주어진다. 그렇다면 민수가 어떤 카드를 낼지 출력하라. 단, 민수가 카드를 내지 못하는 경우는 없다고 가정한다.

 

입력

첫째 줄에 세 개의 자연수 N, M, K가 주어진다. (1 ≤ M ≤ N ≤ 4,000,000, 1 ≤ K ≤ min(M, 10,000))

다음 줄에 카드의 번호를 나타내는 M개의 자연수가 주어진다. 각각의 수들은 1 이상이고 N 이하이며 서로 다르다.

다음 줄에 K개의 자연수가 주어진다. i번째 수는 철수가 i번째로 내는 카드의 번호이다. 철수가 내는 카드 역시 1 이상 N 이하이다.

 

출력

K줄에 걸쳐서 수를 출력한다. i번째 줄에는 민수가 i번째로 낼 카드의 번호가 출력되어야 한다.

---

 

풀이

처음에는 다음과 같은 코드를 제출했다.

from bisect import bisect_right

N, M, K = map(int, input().split())
arr = list(map(int, input().split()))
arr.sort()
order = list(map(int, input().split()))
for i in order:
    idx = bisect_right(arr, i)
    print(arr.pop(idx))

 

아니나 다를까 TLE를 받았다.

arr.pop(idx)의 시간 복잡도가 O(N)이므로 거기서 TLE를 받은 것이 틀림없다.

 

그런데 태그에 '분리 집합'이 있는 것을 보고 뭔가 새로운 아이디어를 떠올려야 했음을 깨달았다. 

 

아이디어는 다음과 같다.

 

1. 먼저 주어진 카드를 정렬한다(이분탐색할거라서)

arr

index	0	1	2	3	4	5	6
value	2	3	4	5	7	8	9

 

2. parent 배열을 만드는데, 이 배열에는 무엇이 저장되어있냐 하면 앞서 정렬한 저 카드의 인덱스들이 저장되어있다. 처음에는 모두 자기 자신을 가리키고 있다.

parent

index	0	1	2	3	4	5	6
value	0	1	2	3	4	5	6

 

3. query가 하나 주어질 때마다, 우리는 민수가 낼 카드를 최대한 빨리 찾아야 한다. 일단 앞서 정렬한 배열(arr)에서 bisect_right를 진행해서 문제의 조건에 맞게 query보다 크면서 가지고 있는 카드 중 가장 작은 카드의 인덱스를 찾을 수 있다. 예를 들어 4가 query로 주어졌다고 하자. 반환하는 인덱스는 3가 될 것이다.

 

3-1. 그러면 이 인덱스 3를 parent 에서 find 할 것이다. 현재는 자기 자신을 가리키고 있으므로 3을 반환한다. 이 find한 값(인덱스)를 d라고 하고, arr[d]를 출력해준다. 답은 5이다.

 

3-2. 5를 썼으므로 이제 더이상 5를 쓰면 안된다. 그러면 우리는 union(d, d+1)을 해준다. 앞으로 또 4와 같은 수들이 주어졌을 때 3을 반환하는 것이 아니라 그 다음으로 큰 수를 반환하는 것이다.

parent

index	0	1	2	3	4	5	6
value	0	1	2	4	4	5	6

 

이를 계속 반복한다.

 

from bisect import *

N, M, K = map(int, input().split())
arr = list(map(int, input().split()))
arr.sort()
order = list(map(int, input().split()))

parent = [i for i in range(M+1)]

def find(x):
    global parent
    if x == parent[x]:
        return x
    parent[x] = find(parent[x])
    return parent[x]

def union(a, b):
    global parent
    a = find(a)
    b = find(b)
    parent[a] = b

for i in order:
    d = find(bisect_right(arr, i))
    print(arr[d])
    union(d, d+1)