PrefixSum 갱신 빠르게 하기

문제

 

2268번: 수들의 합 7

 

---

 

재정의

PrefixSum(구간 합)과 갱신을 하는 단순한 문제.

N이 10**6이고, M이 10**6이다. O(n*log n)으로 풀어야 한다.

M번 만큼의 loop는 무조건 돌아야 하므로,

PrefixSum, Update를 O(log n)이하로 풀어야 한다.

 

누적합을 이용하면 O(1)로 PrefixSum을 구할 수 있다.

하지만 update하는데 O(n)이 들기 때문에, 누적합을 활용할 수 없다.

 

O(log n)으로 PrefixSum과 Update를 하기 위해서는 SegmentTree 또는 FenwickTree를 사용한다.

 

---

SegmentTree

Tree 구현 시에, 주로 배열로 구현하게 된다.

완전 이진 트리를 고려하면, 매 높이 마다 노드 수가 2배씩 증가하기 때문에

1 + 2 + 4 + ... + arr의 길이가 총 필요한 노드의 수가 된다. (2 ** 트리의 높이)만큼 이고,

높이 별로 노드 수가 2배이기 때문에 log를 이용하여 다음과 같이 표현 수 있다.

 

self.segment = [0] * (2 ** ceil(log2(len(arr))+1))
간단하게 4n으로 사용하기도 한다.

 

Init

def init(node, start, end):
    if start == end:
        tree[node] = arr[start]
        return

    mid = (start + end) // 2
    init(2 * node, start, mid)
    init(2 * node + 1, mid + 1, end)
    tree[node] = tree[2 * node] + tree[2 * node + 1]

init(1, 0, N-1)

 

PrefixSum

def query(node, start, end, left, right):
    if right < start or end < left:
        return 0

    if left <= start and end <= right:
        return tree[node]

    mid = (start + end) // 2
    q1 = query(2 * node, start, mid, left, right)
    q2 = query(2 * node + 1, mid + 1, end, left, right)
    return q1 + q2

 

Update

def update(node, start, end, idx, diff):
    if idx < start or end < idx:
        return

    tree[node] += diff

    if start != end:
        mid = (start + end) // 2
        update(2 * node, start, mid, idx, diff)
        update(2 * node + 1, mid + 1, end, idx, diff)

 

---

 

FanwickTree

이 문제는 SegmentTree로 구현하면 아슬아슬하다.

그래서 좀 더 compact한 FanwickTree가 더욱 좋다.

이는 트리의 형태를 만들지 않고, 비트를 활용해서 저장하기 때문이다.

Init

for i in range(1, N+1):
    x = int(input())
    arr[i] = x
    update(i, x)

 

PrefixSum

def prefix_sum(idx: int) -> int:
    res = 0

    while idx > 0:
        res += fenwick_tree[idx]
        idx -= -idx&idx
    return res


def get_sum(st_idx: int, ed_idx: int) -> int:
    return prefix_sum(ed_idx) - prefix_sum(st_idx-1)

 

Update

def update(idx: int, diff: int) -> None:
    while idx <= n:
        fenwick_tree[idx] += diff
        idx += -idx&idx