[백준 2243번] 사탕상자 (Python3)

import sys
input = sys.stdin.readline
M = 2**20

def updatefenwick(i,v):
  while i<=M:
    fenwick[i] += v
    i += i&-i

def sumfenwick(i):
  SUM = 0
  while i:
    SUM += fenwick[i]
    i -= i&-i
  return SUM

def find(x,i,c):
  s = sumfenwick(i)
  if s >= x and s-data[i] < x:
    return i
  if s < x:
    return find(x,i+2**c,c-1)
  return find(x,i-2**c,c-1)

N = int(input())
data,fenwick = [0]*(M+1),[0]*(M+1)

for _ in range(N):
  Q = [*map(int,input().split())]
  if Q[0]==1:
    idx = find(Q[1],2**19,18)
    data[idx] -= 1
    updatefenwick(idx,-1)
    print(idx)
  else:
    data[Q[1]] += Q[2]
    updatefenwick(Q[1],Q[2])

이분탐색과 팬윅트리로 해결하였다.

누적합을 구해야하는 문제이기 때문에 세그먼트 트리보다 팬윅트리가 더 적합하다고 판단하였다.

 

알고리즘을 설명하면,

1. 사탕의 개수에 대한 팬윅트리를 만든다.

2. x번째 사탕을 골라야하는 경우, 이분탐색을 통해 i까지의 누적합이 x이면서 i까지의 누적합-data[i]의 값이 x보다 작아지는 경우를 찾아 i를 출력한다.

 

 

오늘의 교훈) 누적합을 이용해야하는 문제는 팬윅트리를 사용하자