[백준/2912] 백설공주와 난쟁이 (Python)

Problem : https://www.acmicpc.net/problem/2912

2912번: 백설공주와 난쟁이

 

Difficulty : Platinum 2

 

Status : Solved(pypy3)

 

Time : 00:23:41

 

---

 

문제 설명

 

 

백설 공주와 난쟁이 N명과 함께 숲 속에 살고 있다. 난쟁이는 매일 광산에 일하러가고, 백설 공주는 그동안 페이스북을 하고 있다.

매일 아침 난쟁이는 한 줄로 휘파람을 불면서 광산으로 출근을 한다. 백설 공주는 그 주변을 돌아다니면서 난쟁이들 사진을 찍는다.

난쟁이가 광산에 들어가면, 백설 공주는 다시 집으로 돌아간다. 집으로 돌아가면서 찍은 사진 중에 페이스북에 올릴 예쁜 사진을 고른다. 각 난쟁이는 모두 모자를 쓰고 있다. 모자의 색상은 총 C가지가 있다. 사진에 찍힌 난쟁이가 쓰고 있는 모자의 색상 중 절반보다 많은 색이 같은 색이라면 예쁜 사진이다. 즉, 사진에 난쟁이가 K명 찍혀있고, K/2보다 많은 난쟁이의 모자 색이 같다면 예쁜 사진이다.

백설공주가 찍은 사진 M개와 각 사진에 찍힌 난쟁이가 주어졌을 때, 예쁜 사진인지 아닌지를 알아내는 프로그램을 작성하시오.

 

 

입력 및 출력

 

입력

 

첫째 줄에 난쟁이의 수 N과 모자 색상의 수 C가 주어진다. (3 ≤ N ≤ 300,000, 1 ≤ C ≤ 10,000)

둘째 줄에는 각 난쟁이가 쓰고 있는 모자의 색상이 줄을 서 있는 순서대로 주어진다. 색상은 C이하의 자연수로 나타낸다.

셋째 줄에는 사진의 수 M이 주어진다. (1 ≤ M ≤ 10,000)

다음 M개 줄에는 두 정수 A와 B가 주어진다. (1 ≤ A ≤ B ≤ N) 이 줄은 사진의 정보를 의미하고, A번째 난쟁이부터 B번째 난쟁이까지 사진에 찍혔다는 뜻이다.

 

출력

 

출력은 총 M 줄이다. 각 사진이 예쁘지 않다면 "no"를 출력하고, 예쁘다면 "yes X"를 출력한다. 예쁜 사진인 경우에 X는 사진에 절반이 넘는 모자의 색상이다.

 

입력 예시

 

10 3
1 2 1 2 1 2 3 2 3 3
8
1 2
1 3
1 4
1 5
2 5
2 6
6 9
7 10

 

출력 예시

 

no
yes 1
no
yes 1
no
yes 2
no
yes 3

 

 

---

 

풀이

 

구간에 대한 처리를 오프라인으로 진행할 수 있다는 점에 착안하여 오프라인 쿼리 + Mo's로 먼저 도전해 보았다. 이 경우, 쿼리 처리에 대한 시간복잡도는 O((M+N)sqrt(N)) 이 되리라는 가설에서였다. 문제는 절반이 넘는 숫자를 판단해야하므로 C개 값에 대한 체크가 필요했고, 이 체크에 소요되는 시간은 최악의 경우 C가 된다. 즉 실제로는 O((MC+N)sqrt(N))의 시간복잡도로 수행하게 되어 python3에서는 시간 초과, pypy3으로 간신히 통과할 수 있었다.

 

풀이 코드 1 (pypy3)

from collections import defaultdict
import sys
input = sys.stdin.readline

N, C = map(int, input().split())
hats = list(map(int, input().split()))
M = int(input())

hat_dict = defaultdict(int)
ans = [0]*M
queries = [[i] + list(map(lambda x : int(x)-1, input().split())) for i in range(M)]
queries.sort(key = lambda x : (x[1] // (N ** 0.5), x[2]))

def check(K) :
  for key, val in hat_dict.items() :
    if val > K / 2 :
      return key
  return 0

l, r = 0, -1
for idx, i, j in queries :
  while r < j :
    r += 1
    hat_dict[hats[r]] += 1
  while r > j :
    hat_dict[hats[r]] -= 1
    if not hat_dict[hats[r]] :
      del hat_dict[hats[r]]
    r -= 1
  while l < i :
    hat_dict[hats[l]] -= 1
    if not hat_dict[hats[l]] :
      del hat_dict[hats[l]]
    l += 1
  while l > i :
    l -= 1
    hat_dict[hats[l]] += 1
  ans[idx] = check(j - i + 1)

for v in ans :
  if not v :
    print("no")
  else :
    print("yes", v)

반쪽짜리 풀이 완료!

 

(어쨌든 맞았으니) 풀이에 사용할 수 있는 알고리즘 분류를 보자, 무작위화 알고리즘을 발견할 수 있었다. 진짜 충격적인 발상의 전환이기는 했다...

 

우선, 구간의 난쟁이가 과반수 이상인 경우를 생각해 보자. 이 구간에서 임의의 난쟁이를 뽑았을 때, 과반에 속하는 모자일 경우는 1/2 이상이다. 즉 과반에 속하지 않은 모자를 뽑았을 경우는 1/2 이하가 되며, K번 이를 반복할 경우 그 확률은 (1/2)^K 이하로 줄어든다.

 

즉 K번 동안

 

• 유니폼 랜덤하게 모자를 뽑아서
• 그 모자가 구간 내에 과반이라면 무조건 yes.
• 만약 그렇지 않다면 continue

로 K번을 수행함으로써 효과적으로 계산해 볼 수 있는 것이다! 충분한 K값이면 틀릴 확률이 매우 적어진다는 사실을 알 수 있으므로(이를테면 K = 20일 때 2^20 ~= 10^6이므로 1e-6의 오차확률 내), 충분히 적용 가능해보인다(물론 거진 로또맞을 확률로 틀릴 수 있다는 점은 부정할 수 없다).

 

이 경우 과반을 빠르게 계산하려면 전처리를 통하여 O(logN)의 시간 내에 계산 가능하다. 오른쪽 구간 경계의 upper bound, 왼쪽 경계의 lower bound를 구하는 방식으로 적용해볼 수 있다. 즉 총 시간복잡도는 O(MKlogN)이 된다. 오늘도 새롭게 한 수 배워간다...!

 

이 분 블로그 글이 많이 도움이 되었다.

 

https://justicehui.github.io/hard-algorithm/2020/03/23/random/

랜덤으로 문제 풀기

 

풀이 코드 2 (python)

from random import randint
import sys
input = sys.stdin.readline

MAX = 20
N, C = map(int, input().split())
hats = list(map(int, input().split()))
M = int(input())

hatidx = [list() for _ in range(C+1)]
for i in range(N) :
  hatidx[hats[i]].append(i)

def cal(l, r, key) :
  lstart = rstart = 0
  lend = rend = len(hatidx[key])
  while lstart < lend :
    mid = (lstart + lend) // 2
    if hatidx[key][mid] < l :
      lstart = mid + 1
    else :
      lend = mid
  while rstart < rend :
    mid = (rstart + rend) // 2
    if hatidx[key][mid] <= r :
      rstart = mid + 1
    else :
      rend = mid
  return rend - lend

for _ in range(M) :
  l, r = map(lambda x : int(x)-1, input().split())
  K = r - l + 1
  flg = False
  for _ in range(MAX) :
    key = hats[randint(l, r)]
    if cal(l, r, key) > K / 2 :
      flg = True
      break
  if flg :
    print("yes", key)
  else :
    print("no")

풀이 완료!