모독

• A[i] = 생명점이 i인 하수인의 갯수 로 정의되는 배열 A를 생각하자.
• 구해야 하는 것은, A[i]=0인 i의 최소값을 x라고 할때, sum(A[1:x])를 구하는 것이고, 이것은 세그먼트 트리를 이용해서 쉽게 구현 가능하다.
  • 각 구간 노드가 sum과 is_full의 두가지 값을 저장하고 있으면 된다. sum은 구간의 시작 위치부터 가장 처음으로 A[i]가 0인 위치까지의 A[i]의 합이다. is_full은 구간내에 A[i]의 값이 0인 i가 존재하지 않을때 True를 갖는다.
  • 그러면 왼쪽 구간 l과 오른쪽 구간 r을 합쳐서 큰 구간 u에 대해서 저 값을 계산할 때에는, u.is_full은 당연히 l.is_full && r.is_full 이 되고, u.sum은 l.is_full이면 l.sum+r.sum, 그렇지 않으면 그냥 l.sum이 된다.
  • 이렇게 해서 쿼리가 들어올때마다, A[i]를 갱신시켜서 세그먼트 트리를 업데이트하고, 전체 구간에 대한 sum값을 출력하면 된다.
  • 그러나.. 이것은 O(n+qlogn)의 효율적인 시간복잡도를 갖는 풀이임에도 불구하고, 시간 초과에 걸린다. PyPy를 써도 시간 초과. 이것은 문제가 제한 시간이 “3초 (추가시간 없음)“으로 설정되어있는, python에 대한 배려가 없는 문제이기 때문이다..ㅜㅜ
• 시간 복잡도 자체를 줄일수는 없고, 상수라도 줄일수 있는 다른 방법이 필요하다. A[i]=0인 i의 최소값 x를 다른 방식으로 구하고, sum(A[1:x])를 구하는 방법으로 접근하자.
• 이는 A[i]=0인 i를 모두 min heap에 저장하는 방식으로 해결 가능하다. 그리고 A[i]가 0에서 1로 바뀌면 힙에서 i를 빼주고, A[i]가 1에서 0으로 바뀌면 힙에 i를 추가해주면 된다.
• 그러면 쿼리가 들어올때마다 해야할 일이, “튜플을 노드로 갖는 세그먼트 트리 갱신” 에서, “힙에 삽입or삭제 + 펜윅트리로 구간합 계산” 으로 바뀌었다. 둘다 O(logn)에 처리되는 작업이므로 전체 시간 복잡도는 여전히 O(n+qlogn)이지만, 이렇게 바꾸고 나니 PyPy로 약 2.8초에 아슬아슬하게 통과가 가능했다.
• c++로 짠 다른 상위권 코드들을 보면, 힙 대신에, a[i]가 0이면 0을 갖고, 0보다 크면 1을 갖는 펜윅트리를 추가로 만들어서 그것으로 A[i]=0인 i의 최소값 x을 구해서 처리하는 코드들이 있었다. 흠.. 힙을 쓰는 것보다 느릴꺼 같은데;

"""Solution code for "BOJ 16221. 모독".

- Problem link: https://www.acmicpc.net/problem/16221
- Solution link: http://www.teferi.net/ps/problems/boj/16221
"""

import heapq
import sys
from teflib import fenwicktree

MAX_K = 1_000_000


def main():
    Q = int(sys.stdin.readline())

    fenwick = fenwicktree.FenwickTree(MAX_K + 2)
    zeros = list(range(1, MAX_K + 2))
    heapq.heapify(zeros)
    for _ in range(Q):
        T, K = [int(x) for x in sys.stdin.readline().split()]
        if T == 1:
            fenwick.update(K, 1)
            if fenwick.get(K) == 1 and K == zeros[0]:
                while fenwick.get(zeros[0]) > 0:
                    heapq.heappop(zeros)
        else:
            fenwick.update(K, -1)
            if fenwick.get(K) == 0:
                heapq.heappush(zeros, K)
        print(fenwick.query(1, zeros[0]))


if __name__ == '__main__':
    main()

• Dependency: teflib.fenwicktree.FenwickTree