[백준/16978] 수열과 쿼리 22 (Python)

Problem : https://www.acmicpc.net/problem/16978

16978번: 수열과 쿼리 22

 

Difficulty : Platinum 4

 

Status : Solved

 

Time : 00:26:23

 

---

 

문제 설명

 

 

길이가 N인 수열 A1, A2, ..., AN이 주어진다. 이때, 다음 쿼리를 수행하는 프로그램을 작성하시오.

• 1 i v: Ai = v로 변경한다.
• 2 k i j: k번째 1번 쿼리까지 적용되었을 때, Ai, Ai+1, ..., Aj의 합을 출력한다.

 

 

입력 및 출력

 

입력

 

첫째 줄에 수열의 크기 N (1 ≤ N ≤ 100,000)이 주어진다.
둘째 줄에는 A1, A2, ..., AN이 주어진다. (1 ≤ Ai ≤ 1,000,000)
셋째 줄에는 쿼리의 개수 M (1 ≤ M ≤ 100,000)이 주어진다.
넷째 줄부터 M개의 줄에는 쿼리가 한 줄에 하나씩 주어진다. 1번 쿼리의 경우 1 ≤ i ≤ N, 1 ≤ v ≤ 1,000,000 이고, 2번 쿼리의 경우 1 ≤ i ≤ j ≤ N이고, 0 ≤ k ≤ (쿼리가 주어진 시점까지 있었던 1번 쿼리의 수)이다.
입력으로 주어지는 모든 수는 정수이다.

 

출력

 

모든 2번 쿼리마다 합을 출력한다.

 

입력 예시

 

5
1 2 3 4 5
7
1 2 5
2 0 1 3
2 1 1 3
1 4 2
2 0 2 5
2 1 2 5
2 2 2 5

 

출력 예시

 

6
9
14
17
15

 

 

---

 

풀이

 

1번 쿼리 및 2번 쿼리의 특성상 세그먼트 트리를 떠올리는 것까지는 어떻게든 생각할 수 있을 것이다. 문제는, 이 2번 쿼리가 특정 시점의 1번 쿼리에서의 부분합을 요구한다는 점이다. 즉 이 경우 순차적으로 쿼리를 처리하려면 1번 쿼리의 개수만큼 세그먼트 트리의 복사본이 필요하며, 이는 심각한 메모리상의 비효율성을 초래한다.

 

즉 다른 방법으로 접근해 볼 필요가 있다. 우리가 쿼리 순서를 조작할 수 있다면, 1번 쿼리를 계속해서 처리하다가, 조건이 만족되는 2번 쿼리를 먼저 처리하는 방식을 사용해 볼 법 하다. 즉 오프라인 쿼리 문제로 볼 수 있다. 2번 쿼리를 등장 순으로 저장하고, 요구되는 1번 쿼리 햇수를 기준으로 정렬한 후, 1번 쿼리의 처리 횟수에 따라 정확히 일치하는 2번 쿼리를 처리하면 되는 방식이다.

 

풀이 코드

from collections import deque
import sys
input = sys.stdin.readline

N = int(input())
num_list = list(map(int, input().split()))
M = int(input())
search_q, update_q = list(), list()
sidx, uidx = 0, 0

for i in range(M) :
  q, *cmd = map(int, input().split())
  if q == 2 :
    search_q.append([sidx] + cmd)
    sidx += 1
  else :
    update_q.append([uidx] + cmd)
    uidx += 1

class SegTree :
  def __init__(self) :
    self.tree = [0]*(4*N)
    def _init(start, end, idx) :
      if start == end :
        self.tree[idx] = num_list[start]
        return
      mid = (start + end) // 2
      _init(start, mid, idx*2)
      _init(mid+1, end, idx*2+1)
      self.tree[idx] = self.tree[idx*2] + self.tree[idx*2+1]
    _init(0, N-1, 1)

  def update(self, target, val) :
    def _update(start, end, idx) :
      if target < start or end < target :
        return
      if start == end :
        self.tree[idx] = val
        return
      mid = (start + end) // 2
      _update(start, mid, idx*2)
      _update(mid+1, end, idx*2+1)
      self.tree[idx] = self.tree[idx*2] + self.tree[idx*2+1]
    _update(0, N-1, 1)

  def search(self, left, right) :
    def _search(start, end, idx) :
      if right < start or end < left :
        return 0
      if left <= start <= end <= right :
        return self.tree[idx]
      mid = (start + end) // 2
      l = _search(start, mid, idx*2)
      r = _search(mid+1, end, idx*2+1)
      return l + r
    return _search(0, N-1, 1)

search_q = deque(sorted(search_q, key = lambda x : x[1]))
ans = [0]*sidx
segtree = SegTree()
for uidx, idx, val in update_q :
  while search_q and search_q[0][1] == uidx :
    sidx, _, left, right = search_q.popleft()
    ans[sidx] = segtree.search(left-1, right-1)
  segtree.update(idx-1, val)
while search_q :
  sidx, _, left, right = search_q.popleft()
  ans[sidx] = segtree.search(left-1, right-1)

print(*ans, sep = '\n')

풀이 완료!