파이썬으로 풀어보는 백준 2912번: 백설공주와 난쟁이

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

2912번: 백설공주와 난쟁이

 

 

내 풀이(1):

import copy
import sys
input = sys.stdin.readline
N, C = map(int, input().split())
data = list(map(int, input().split()))
M = int(input())
 
record = [0] * (N+1)
record[0] = {}
rec = {}
for idx, d in enumerate(data):
    if d in rec:
        rec[d] += 1
    else:
        rec[d] = 1
    record[idx + 1] = copy.deepcopy(rec)
 
for _ in range(M):
    A, B = map(int, input().split())
    for num in record[B]:
        if num in record[A-1]:
            num_A = record[A-1][num]
        else:
            num_A = 0
        num_B = record[B][num]
        if num_B - num_A > (B-A+1) / 2:
            print('yes', num)
            break
    else:
        print('no')

python3, 시간 초과

처음으로 시도한 방법은 각 단계까지의 모자 색깔에 대한 기록을 딕셔너리의 형태로 각각 남기는 것이었다. 각 단계에 서로 다른 딕셔너리를 남겨야 하기 때문에 deepcopy를 이용해주었고 deepcopy 과정에서 원소의 개수만큼의 시간을 소요하게 된다. 또한 A, B라는 두 단계가 선택되면 두 단계에 대한 특정 색깔의 차이를 계산할 때 최대 10000가지 색에 대해 검토하게 된다. deepcopy 과정과 10000가지 색에 대한 검토에 의해 시간 초과가 난 것으로 판단된다.

 

 

내 풀이(2):

 

import sys
input = sys.stdin.readline
 
def cal(start, end, idx, A, B):
    if A <= start and end <= B:
        for num in segmentTree[idx]:
            if num in result:
                result[num] += segmentTree[idx][num]
            else:
                result[num] = segmentTree[idx][num]
        return
 
    if end < A or B < start:
        return
 
    mid = (start + end) // 2
    cal(start, mid, idx * 2, A, B)
    cal(mid + 1, end, idx * 2 + 1, A, B)
    return
 
 
def make(start, end, idx):
    record = {}
    if start == end:
        record[data[start - 1]] = 1
        segmentTree[idx] = record
        return record
 
    mid = (start + end) // 2
    left = make(start, mid, idx * 2)
    right = make(mid + 1, end, idx * 2 + 1)
 
    for num in left:
        if num in record:
            record[num] += left[num]
        else:
            record[num] = left[num]
 
    for num in right:
        if num in record:
            record[num] += right[num]
        else:
            record[num] = right[num]
 
    segmentTree[idx] = record
    return record
 
 
N, C = map(int, input().split())
data = list(map(int, input().split()))
M = int(input())
segmentTree = [{} for _ in range(1048562)]
make(1, N, 1)
for _ in range(M):
    A, B = map(int, input().split())
    result = {}
    cal(1, N, 1, A, B)
    check = (B - A + 1) // 2
    for num in result:
        if result[num] > check:
            print('yes', num)
            break
    else:
        print('no')

import sys
import random
 
input = sys.stdin.readline
def cal(start, end, idx, A, B, num):
    if A <= start and end <= B:
        if num in segmentTree[idx]:
            return segmentTree[idx][num]
        else:
            return 0
 
    if end < A or B < start:
        return 0
 
    res = 0
    mid = (start + end) // 2
    res += cal(start, mid, idx * 2, A, B, num)
    res += cal(mid + 1, end, idx * 2 + 1, A, B, num)
    return res
 
 
def make(start, end, idx):
    record = {}
    if start == end:
        record[data[start - 1]] = 1
        segmentTree[idx] = record
        return record
 
    mid = (start + end) // 2
    left = make(start, mid, idx * 2)
    right = make(mid + 1, end, idx * 2 + 1)
 
    for num in left:
        if num in record:
            record[num] += left[num]
        else:
            record[num] = left[num]
 
    for num in right:
        if num in record:
            record[num] += right[num]
        else:
            record[num] = right[num]
 
    segmentTree[idx] = record
    return record
 
 
N, C = map(int, input().split())
data = list(map(int, input().split()))
M = int(input())
segmentTree = [{} for _ in range(1050000)]
make(1, N, 1)
for _ in range(M):
    A, B = map(int, input().split())
    check = (B - A + 1) // 2
 
    check_color = {}
    for _ in range(20):
        random_color = data[random.randint(A, B)-1]
        if random_color in check_color:
            continue
        else:
            check_color[random_color] = 1
        if cal(1, N, 1, A, B, random_color) > check:
            print('yes', random_color)
            break
    else:
        print('no')

python3, 4136ms

10000가지나 되는 색상에 대해 모두 조사해보는 것은 많은 시간을 요구한다. 때문에 여기서 수학적, 확률적 접근이 이용된다. 특정 구간에 과반 수 이상의 특정 색깔이 존재한다고 가정하면 임의로 한 명의 난쟁이를 뽑아 모자 색을 확인했을 때, 그 색이 과반수가 아닐 확률은 50%이다. 이어서 두 명째를 뽑아 연속으로 과반수가 아닐 확률은 50% * 50% = 25% 이다. 이와 같은 계산을 해볼 때, 20명의 난쟁이를 확인하여 과반수가 넘는 모자의 색이 존재함에도 불구하고 그 색이 단 한 번도 뽑히지 않을 확률은 0.000095367% 로 매우 희박하다. 따라서 10000가지 색이 아닌 20명의 난쟁이를 뽑아 검토하는 것으로 문제의 조건을 충분히 해결할 수 있다. 위 알고리즘들을 이용하여 4136ms의 결과로 문제를 통과할 수 있었다.

 

 

내 풀이(4):

import random
import sys
input = sys.stdin.readline
def cal_cnt(random_color, A, B):
    low = 0
    high = len(record[random_color]) - 1
 
    # Lower_bound
    while low < high:
        m = (low + high) // 2
        if record[random_color][m] < A:
            low = m + 1
        else:
            high = m
    if record[random_color][high] < A:
        return 0
    else:
        start = high
 
    # upper_bound
    low = 0
    high = len(record[random_color])
    while low < high:
        m = (low + high) // 2
        if record[random_color][m] <= B:
            low = m + 1
        else:
            high = m
    high -= 1
    if record[random_color][high] > B:
        return 0
    else:
        end = high
 
    return end - start + 1
 
 
N, C = map(int, input().split())
caps = [0] + list(map(int, input().split()))
record = {}
for idx, color in enumerate(caps):
    if color not in record:
        record[color] = [idx]
    else:
        record[color].append(idx)
 
M = int(input())
visited = {}
for _ in range(M):
    A, B = map(int, input().split())
    check = (B - A + 1) // 2
    check_color = {}
    if A in visited:
        if B in visited[A]:
            print(visited[A][B])
            continue
 
    for _ in range(20):
        random_color = caps[random.randint(A, B)]
        if random_color in check_color:
            continue
        else:
            check_color[random_color] = 1
        cnt = cal_cnt(random_color, A, B)
        if cnt > check:
            print('yes', random_color)
            if A not in visited:
                visited[A] = {}
            visited[A][B] = 'yes ' + str(random_color)
            break
    else:
        print('no')
        if A not in visited:
            visited[A] = {}
        visited[A][B] = 'no'

python3, 1068ms

segment tree를 만드는 데 시간이 오래 걸리지 않을까 하는 생각에 특정 구간에서 특정 색에 대한 모자 수를 세는 방법을 upper_bound와 lower_bound를 이용하여 해결하였다. 이를 이용하기 위해 미리 각 모자 색마다 등장하는 index를 저장해주고 난수로 발생시킨 random_color에 대한 리스트에서 A, B 조건에 만족하는 lower_bound와 upper_bound의 index 차이로 이 역시 O(logN)이라는 시간복잡도로 개수를 구할 수 있다. 다만 이 방법은 처음 데이터를 읽어드리는 순간 바로 색깔별 리스트를 만들어줄 수 있기 때문에 사전 준비에 시간을 필요로 하지 않는다.

 

Python 언어로는 푼 사람도 없고 인터넷으로도 풀이를 찾아볼 수 없어 문제를 푸는 데 여러 알고리즘에 대해 고민하고 공부하여 어렵게 문제를 해결하였다. 효율성에 관한 알고리즘 해결에 힘을 쓰고 있지만 여전히 배워야하고 실수하지 않아야 할 부분이 많은 것 같다.