class LazySegTree:
def __init__(self, arr):
self.n = len(arr)
self.size = 1
while self.size < self.n:
self.size <<= 1
self.tree = [0] * (self.size * 2)
self.lazy = [0] * (self.size * 2)
self.length = [1] * (self.size * 2)
# 각 노드가 담당하는 구간 길이
for i in range(self.size - 1, 0, -1):
self.length[i] = self.length[i * 2] + self.length[i * 2 + 1]
# 리프 초기화
for i in range(self.n):
self.tree[self.size + i] = arr[i]
# 내부 노드 구성
for i in range(self.size - 1, 0, -1):
self.pull(i)
def pull(self, node):
self.tree[node] = self.tree[node * 2] + self.tree[node * 2 + 1]
def apply(self, node, value):
"""
node 전체 구간에 value를 더한다.
"""
self.tree[node] += value * self.length[node]
self.lazy[node] += value
def push(self, node):
"""
node에 쌓인 lazy 값을 자식에게 내린다.
"""
if self.lazy[node] != 0:
value = self.lazy[node]
self.apply(node * 2, value)
self.apply(node * 2 + 1, value)
self.lazy[node] = 0
def push_path(self, node):
"""
루트에서 node까지 내려가며 lazy를 전파한다.
쿼리나 업데이트 전에 경계 노드의 lazy를 정확히 반영하기 위해 사용한다.
"""
stack = []
node >>= 1
while node:
stack.append(node)
node >>= 1
for x in reversed(stack):
self.push(x)
def rebuild_path(self, node):
"""
node에서 루트까지 올라가며 값을 다시 계산한다.
업데이트 후 사용한다.
"""
node >>= 1
while node:
self.pull(node)
self.tree[node] += self.lazy[node] * self.length[node]
node >>= 1
def range_add(self, left, right, value):
"""
[left, right)에 value를 더한다.
"""
left += self.size
right += self.size
l0, r0 = left, right
# 경계까지 lazy를 미리 내려줌
self.push_path(l0)
self.push_path(r0 - 1)
while left < right:
if left & 1:
self.apply(left, value)
left += 1
if right & 1:
right -= 1
self.apply(right, value)
left >>= 1
right >>= 1
# 바뀐 경로만 다시 계산
self.rebuild_path(l0)
self.rebuild_path(r0 - 1)
def range_sum(self, left, right):
"""
[left, right)의 합을 구한다.
"""
left += self.size
right += self.size
# 경계까지 lazy를 내려줘야 정확한 값을 얻을 수 있음
self.push_path(left)
self.push_path(right - 1)
result = 0
while left < right:
if left & 1:
result += self.tree[left]
left += 1
if right & 1:
right -= 1
result += self.tree[right]
left >>= 1
right >>= 1
return result
<내가 다듬음> (= 사실 똑같음) (= 똑같다 까진 아니고 리프 조건 따져서 push, pull 하도록 바꿈)
class SegTree:
def __init__(self, arr):
self.n = len(arr)
self.size = 1
while self.size < self.n:
self.size <<= 1
self.tree = [0] * self.size * 2
self.lazy = [0] * self.size * 2
self.cnt = [0] * self.size * 2
for i in range(self.n):
self.tree[i+self.size] = arr[i]
self.cnt[i+self.size] = 1
for i in range(self.size-1, 0, -1):
self.pull(i)
self.cnt[i] = self.cnt[i*2] + self.cnt[i*2+1]
def pull(self, node):
if node < self.size:
self.tree[node] = self.tree[node*2] + self.tree[node*2+1]
def apply(self, node, value):
self.tree[node] += value * self.cnt[node]
self.lazy[node] += value
def push(self, node):
if self.lazy[node]:
if node < self.size:
self.apply(node*2, self.lazy[node])
self.apply(node*2+1, self.lazy[node])
self.lazy[node] = 0
def push_path(self, node):
stack = []
node >>= 1
while node:
stack.append(node)
node >>= 1
for x in reversed(stack):
self.push(x)
def build(self, node):
node >>= 1
while node:
self.pull(node)
self.tree[node] += self.lazy[node] * self.cnt[node]
node >>= 1
def update(self, left, right, value):
left += self.size
right += self.size
l, r = left, right
self.push_path(l)
self.push_path(r-1)
while left < right:
if left & 1:
self.apply(left, value)
left += 1
if right & 1:
right -= 1
self.apply(right, value)
left >>= 1
right >>= 1
self.build(l)
self.build(r-1)
def query(self, left, right):
left += self.size
right += self.size
self.push_path(left)
self.push_path(right-1)
ret = 0
while left < right:
if left & 1:
ret += self.tree[left]
left += 1
if right & 1:
right -= 1
ret += self.tree[right]
left >>= 1
right >>= 1
return ret