[백준 2042번] 구간 합 구하기 (Python3)

import sys
input = sys.stdin.readline

N,M,K = map(int,input().split())

data = [0]*(N+1)
for i in range(N):
  data[i] = int(input())

segtree = {}

def makeseg(start,end,x):
  if start == end:
    segtree[x] = data[start]
  else:
    mid = (start+end)//2
    segtree[x] = makeseg(start,mid,x*2) + makeseg(mid+1,end,x*2+1)
  return segtree[x]

makeseg(0,N-1,1)

def editseg(n,dif,start,end,x):
  segtree[x] += dif
  if start == end == n:
    return
  mid = (start+end)//2
  if n <= mid:
    editseg(n,dif,start,mid,x*2)
  else:
    editseg(n,dif,mid+1,end,x*2+1)

def sumseg(n1,n2,start,end,x):
  if n1 == start and n2 == end:
    return segtree[x]
  mid = (start+end)//2
  if n2 <= mid:
    return sumseg(n1,n2,start,mid,x*2)
  if n1 > mid:
    return sumseg(n1,n2,mid+1,end,x*2+1)  
  return sumseg(n1,mid,start,mid,x*2)+sumseg(mid+1,n2,mid+1,end,x*2+1)

for i in range(M+K):
  a,b,c = map(int,input().split())
  if a == 1:
    dif = c-data[b-1]
    data[b-1] = c
    editseg(b-1,dif,0,N-1,1)
  else:
    print(sumseg(b-1,c-1,0,N-1,1))

 

세그먼트 트리 문제이다. 세그먼트 트리라는게 있다는건 알고 있었는데, 미루고 미루다가 드디어 한번 풀어봤다.

 

최대한 코드는 참고하지 않고 원리만 읽어본 다음에 직접 구현해보려고 노력했다.

원리는 41. 세그먼트 트리(Segment Tree) : 네이버 블로그 (naver.com) 여기에서 참고했다.

 

코드를 다 짜고나서 제출한 후, 다른 사람들 코드와 비교해 보았는데 큰 차이점이 보이지 않았다. 직접 짰는데도 잘 짠 것 같아 만족스러웠다.

 

 

오늘의 교훈) 세그먼트 트리를 응용하는 문제를 풀어보자