[백준 파이썬] 2042 구간 합 구하기

https://www.acmicpc.net/problem/2042

2042번: 구간 합 구하기

문제

구간합을 구하는 문제이다. 다만 중간중간 구간에서 값이 변경된다.

---

 

풀이

dp를 이용하여 구간합을 계산하면 중간중간 값이 바뀌기 때문에 O(N)의 시간 복잡도를 가지게 되고 시간초과가 난다.

따라서 세그먼트 트리를 이용하여 O(log(n))의 시간복잡도로 문제를 해결해야 한다.

 

tree = [0, 15, 6, 9, 3, 3, 4, 5, 1, 2]의 트리는 다음과 같은 구조를 가진다.

이 때 트리의 가장 아래 노드에 인덱스를 지정할 수 있고 각각의 자식의 합을 부모의 값으로 가질 수 있다.

해당 트리를 만들기 위한 코드는 다음과 같다.

def makeTree(x, left, right):
  #구간에 데이터가 1개일경우
  if left == right:
    segtree[x] = num[left]
    return segtree[x]
  
  #구간에 데이터가 여러개일경우
  #부모 노드의 구간을 둘로나눔
  mid = (left + right) // 2
  #자식노드
  leftValue = makeTree(2*x, left, mid)
  rightValue = makeTree(2*x+1, mid+1, right)
  #부모노드는 자식 노드의 합
  segtree[x] = leftValue + rightValue
  return segtree[x]

 

 

이렇게되면 각 범위에 해당하는 합을 구할 수 있다.

예를들어 1~4 (2, 3, 4, 5)의 구간의 합을 구하려면 다음과 같이 구할 수 있다.

포함되는 구간의 가장 상위 노드를 더하면 해당 구간의 합이 된다.

즉 2 + 3 + 9 = 14가 답이다.

#b~c의 구간 합 구하기
#트리의 구간 left~right
#현재 노드 x
def findTree(b, c, x, left, right):
  #1. 구하고 싶은 구간(b~c)이 현재 트리 구간에 포함되지 않는경우
  if c < left or right < b:
    return 0
  
  #2. 구하고 싶은 구간이 현재 트리에 포함될경우
  if b <= left and right <= c:
    return segtree[x]
    
  #3. 구간이 겹치는 경우
  mid = (left + right) // 2
  leftValue = findTree(b, c, x*2, left, mid)
  rightValue = findTree(b, c, x*2+1, mid+1, right)
  return leftValue + rightValue

 

트리의 값을 바꾸기 위해서는 바꿀 노드의 부모노드를 모두 변경해주면 된다.

#트리 변경
def updateTree(x, left, right, idx, val):
  #구간에 데이터가 1개일경우 
  if left == right == idx:
    segtree[x] = val
    return
  #현재 구간에 idx가 포함되지 않는경우
  if idx < left or right < idx:
    return
  #자식 노드에 idx가 포함되면 부모노드도 값이 변경
  mid = (left + right) // 2
  updateTree(x*2, left, mid, idx, val)
  updateTree(x*2+1, mid+1, right, idx, val)
  segtree[x] = segtree[x*2] + segtree[x*2+1]

 

 

---

코드

n, m, k = map(int, input().split())
num = [int(input()) for _ in range(n)]


segtree = [0 for _ in range(4*n)]

def makeTree(x, left, right):
  #구간에 데이터가 1개일경우
  if left == right:
    segtree[x] = num[left]
    return segtree[x]
  
  #구간에 데이터가 여러개일경우
  #부모 노드의 구간을 둘로나눔
  mid = (left + right) // 2
  #자식노드
  leftValue = makeTree(2*x, left, mid)
  rightValue = makeTree(2*x+1, mid+1, right)
  #부모노드는 자식 노드의 합
  segtree[x] = leftValue + rightValue
  return segtree[x]

makeTree(1, 0, n-1)

#b~c의 구간 합 구하기
#트리의 구간 left~right
#현재 노드 x
def findTree(b, c, x, left, right):
  #1. 구하고 싶은 구간(b~c)이 현재 트리 구간에 포함되지 않는경우
  if c < left or right < b:
    return 0
  
  #2. 구하고 싶은 구간이 현재 트리에 포함될경우
  if b <= left and right <= c:
    return segtree[x]
  
  #3. 구간이 겹치는 경우
  mid = (left + right) // 2
  leftValue = findTree(b, c, x*2, left, mid)
  rightValue = findTree(b, c, x*2+1, mid+1, right)
  return leftValue + rightValue

#트리 변경
def updateTree(x, left, right, idx, val):
  #구간에 데이터가 1개일경우 
  if left == right == idx:
    segtree[x] = val
    return
  #현재 구간에 idx가 포함되지 않는경우
  if idx < left or right < idx:
    return
  #자식 노드에 idx가 포함되면 부모노드도 값이 변경
  mid = (left + right) // 2
  updateTree(x*2, left, mid, idx, val)
  updateTree(x*2+1, mid+1, right, idx, val)
  segtree[x] = segtree[x*2] + segtree[x*2+1]


for _ in range(m+k):
  a, b, c = map(int, input().split())
  if a == 1:
    updateTree(1, 0, n-1, b-1, c)
  else:
    s = findTree(b-1, c-1, 1, 0, n-1)
    print(s)