[Python] SegmentTree와 LazyPropagation

🎱 Lazy Propagation

값을 빠르게 업데이트하기 위해 해당 값 사용 전까지 직접 업데이트하지 않고 미뤄두는 최적화 기법이다.

• lazy 배열에 미룬 갱신 값 저장
• propagate 함수로 노드를 사용할 때 갱신 값 업데이트

🎱 코드 구현

⚽️ 기존 Segment Tree

##### non_LazyPropagation
def init_segtree(nums, tree, node, l, r):
    if l == r:
        tree[node] = nums[l]
    else:
        mid = (l + r) // 2
        l_child = init_segtree(nums, tree, node*2, l, mid)
        r_child = init_segtree(nums, tree, node*2+1, mid+1, r)
        tree[node] = l_child + r_child
    return tree[node]

def query_segtree(nums, tree, node, l, r, target_l, target_r):
    if r < target_l or l > target_r:
        return 0
    elif target_l <= l and r <= target_r:
        return tree[node]
    else:
        mid = (l + r) // 2
        l_child = query_segtree(nums, tree, node*2, l, mid, target_l, target_r)
        r_child = query_segtree(nums, tree, node*2+1, mid+1, r, target_l, target_r)
        return l_child + r_child

def update_segtree(nums, tree, node, l, r, target_l, target_r, added_val):
    if r < target_l or l > target_r:
        return
    elif target_l <= l and r <= target_r and l == r:
        tree[node] += added_val
        return
    else:
        mid = (l + r) // 2
        update_segtree(nums, tree, node*2, l, mid, target_l, target_r, added_val)
        update_segtree(nums, tree, node*2+1, mid+1, r, target_l, target_r, added_val)
        tree[node] = tree[node*2] + tree[node*2+1]

기존 Segment Tree와 비교하며 아래의 Lazy Propagation 적용 Segment Tree를 살펴보자.

⚽️ Lazy Propagation 적용 Segment Tree

#### LazyPropagation
def propagate_segtree(lazy, tree, node, l, r):
    if lazy[node] != 0:
        tree[node] += (r - l + 1) * lazy[node]
        if l != r:
            lazy[node*2] += lazy[node]
            lazy[node*2+1] += lazy[node]
        lazy[node] = 0

def query_segtree(lazy, tree, node, l, r, target_l, target_r):
    propagate_segtree(lazy, tree, node, l, r)
    
    if r < target_l or l > target_r:
        return 0
    elif target_l <= l and r <= target_r:
        return tree[node]
    else:
        mid = (l + r) // 2
        l_child = query_segtree(lazy, tree, node*2, l, mid, target_l, target_r)
        r_child = query_segtree(lazy, tree, node*2+1, mid+1, r, target_l, target_r)
        return l_child + r_child

def update_segtree(lazy, tree, node, l, r, target_l, target_r, add_val):
    propagate_segtree(lazy, tree, node, l, r)
    
    if r < target_l or l > target_r:
        return
    elif target_l <= l and r <= target_r:
        lazy[node] += add_val
        return
    else:
        mid = (l + r) // 2
        update_segtree(lazy, tree, node*2, l, mid, target_l, target_r, add_val)
        update_segtree(lazy, tree, node*2+1, mid+1, r, target_l, target_r, add_val)
        tree[node] = tree[node*2] + tree[node*2+1]

Segment Tree에 Lazy Propagation을 적용한 예시를 살펴보자.

⚽️ lazy

lazy 는 각 노드에 대해 미뤄놓았던 계산값을 저장하는 배열이다.

각 노드는 “특정 구간”을 의미하고있기 때문에 이 구간에 해당하는 모든 노드가 저장된 값만큼 업데이트 되어야 함을 의미한다.

⚽️ propagate_segtree

propagate_segtree 에서 수행하는 작업은 현재 노드에 대해 lazy 값이 있다면 노드 값을 업데이트 해준다.
그리고 만약 leaf 노드가 아니라면, 현재 노드의 자식 노드 또한 업데이트 되어야하는 값이었으므로 자식 노드의 lazy 값에 누적해준다.

즉, 여기서는 현재 노드의 미뤄둔 값 갱신만 이루어진다.

⚽️ update_segtree

그럼 update_segtree에서는 무슨 작업이 일어나나?

기존 update_segtree 는 세가지 경우로 나누어 작업을 수행한다.

1. 현재 노드의 담당 범위와 타겟 범위가 완전히 겹치지 않는 경우
   ➡️ 작업을 수행하지 않는다.
2. 현재 노드의 담당 범위와 타겟 범위가 완전히 겹치면서 leaf 노드인 경우
   ➡️ 현재 노드의 tree 값을 업데이트한다.
3. 현재 노드의 담당 범위와 타겟 범위가 일부 겹치는 경우
   ➡️ 자식 노드를 재귀 호출한 뒤, 업데이트된 자식 노드들을 통해 현재 노드를 업데이트한다.

lazy 값은 해당 노드의 모든 자식 노드에게 업데이트 되어야 하는 값을 저장하므로, 2번 경우에만 업데이트할 수 있다.
또한 lazy값을 자식노드로 전파시키는 것은 propagate_segtree 함수의 역할이므로, leaf 노드까지 내려갈 필요 없이 재귀를 중단한다.

이러한 차이점을 반영하여 Lazy Propagation을 적용하면 다음과 같은 흐름을 가지게된다.

1. 모든 update_segtree 수행 시작에 propagate_segtree 를 호출하여 현재 노드의 업데이트된 정확한 값을 사용하도록 한다.
2. 현재 노드의 담당 범위와 타겟 범위가 완전히 겹치지 않는 경우
   ➡️ 작업을 수행하지 않는다.
3. 현재 노드의 담당 범위와 타겟 범위가 완전히 겹치는 경우
   ➡️ 현재 노드의 lazy 값을 업데이트한다.
4. 현재 노드의 담당 범위와 타겟 범위가 일부 겹치는 경우
   ➡️ 자식 노드를 재귀 호출한 뒤, 업데이트된 자식 노드들을 통해 현재 노드를 업데이트한다. (재귀 호출을 통해 자식 노드에서도 똑같은 과정이 일어난다.)

⚽️ query_segtree

query_segtree 에서와 유사하게, 모든 query_segtree 수행 시작에 propagate_segtree 를 호출하여 먼저 현재 노드의 업데이트된 정확한 값을 사용하도록 한다.

나머지는 기존 query_segtree와 똑같다.

🎱 언제 사용할까?

구간에 대한 값 업데이트가 빠르게 이루어져야할 경우 사용할 수 있는 방법이다.

현재 모든 설명이 세그먼트 트리를 기준으로 작성되었으나, 이외에 다양한 상황에서도 사용 가능하다.