[BOJ] 1572 - 중앙값

[백준] 1572 - 중앙값

예전에 풀었다가 틀린 뒤로 한동안 잡지 않았던 문제다.(TMI)
이게 대체 왜 펜윅트리인지 엄청 고민하다가 팍 하고 해결책이 떠오른걸 보아 아무래도 헛공부를 했던 것은 아니었던 걸로..
각설하고, 문제를 풀어보자.

접근

중앙값을 구하려면 우선 주어진 배열을 정렬해야 한다. 하지만 주어진 배열을 정렬하고, 매번 값을 찾아 제거하는 과정은 상당한 시간 복잡도를 필요로 한다.(단순 sort로 구현했을 경우 klogk * N * k(제거에 걸리는 시간)
따라서 부분배열을 일일이 정렬하지 않고 부분 배열에 들어있는 수의 갯수정보를 펜윅트리로 저장하여 중앙값 인덱스에 해당하는 숫자를 이분탐색으로 탐색한다. 이렇게 되면 매 탐색마다 O(log65536)(즉, 상수시간)이 걸리고, 업데이트 하는 과정에서 O(logN)이 걸린다.

소스코드

from collections import deque

tree = [0 for _ in range(65540)]


def update(idx, val):
    while idx <= 65539:
        tree[idx] += val
        idx += idx & (-idx)


def find(idx):
    rtn = 0
    while idx > 0:
        rtn += tree[idx]
        idx -= idx & (-idx)
    return rtn


n, k = map(int, input().split())
dq = deque()


def biSearch():
    L = 1
    R = 65537
    median = (k+1)//2
    while True:
        mid = (L+R)//2
        cnt = find(mid)
        if L == mid:
            if cnt < median:
                return R-1
            return L-1
        if cnt < median:
            L = mid
        else:
            R = mid


ans = 0

for i in range(n):
    dq.append(int(input()))
    update(dq[-1]+1, 1)
    if len(dq) > k:
        update(dq.popleft()+1, -1)
    if len(dq) == k : ans += biSearch()

print(ans)