백준 2042 구간 합 구하기/세그먼트트리

문제로 이동

---

목차

문제 설명

구간합 문제다. 이 게시물에서는 세그먼트 트리를 이용해 해결한다.

충분히 단순 반복문으로 해결할 수 있지만 그런 경우, 구간 합을 구하는데 O(n)이 걸린다. 이런 과정이 M번 반복된다면 O(MN)이 된다. 구간합 뿐만 아니라 특정 위치의 수를 변경시키는 기능도 있다. 해당 위치의 수가 바뀌면 그보다 큰 위치의 합도 모두 변경해야 하므로 TLE가 날 수 밖에 없다.

 

세그먼트 트리를 이용하면 O(Mlog(N))으로 해결할 수 있다.

---

소스 코드

소스코드

import sys
input = sys.stdin.readline

tree = list()
nodes = list()

#ChangeIdx : 수정할 인덱스, ChangeNum : 수정할 숫자
def update(start, end, index, ChangeIdx, ChangeNum):
    #범위를 벗어나는 경우
    if ChangeIdx<start or ChangeIdx>end:return
    #하나씩 내려가면서 다른 원소들도 갱신해야함.
    tree[index] += ChangeNum-nodes[ChangeIdx]#값의 변동치만큼 가감시켜줌
    if start==end: return
    mid = (start+end)//2
    update(start,mid,index*2,ChangeIdx,ChangeNum)
    update(mid+1,end,index*2+1,ChangeIdx,ChangeNum)
    

#합을 출력하는 함수
def query(start,end,index,left,right):
    # 입력된 쿼리의 범위가 트리의 범위를 벗어나는 경우
    if left>end or right<start:
        return 0
    #범위 내인 경우
    if start>=left and end<=right:
        return tree[index]
    mid = (start+end)//2
    return query(start,mid,index*2,left,right)+query(mid+1,end,index*2+1,left,right)
#누적 합으로 트리 구성해줌
def init(start,end,index):
    if start==end:
        tree[index] = nodes[start]
        return tree[index]
    mid = (start+end)//2
    tree[index] = init(start,mid,index*2)+init(mid+1,end,index*2+1)
    return tree[index]

if __name__ == "__main__":
    #노드수, 수의변경 수, 구간합 수
    n,m,k = map(int,input().rstrip().split())
    for _ in range(n):
        nodes.append(int(input()))
    tree = [0]*(n*4)
    init(0,n-1,1)
    for _ in range(m+k):
        a,b,c = map(int,input().rstrip().split())        
        # a==1 : b번째 수를 c로 변경
        # a==2 : b부터 c까지 합을 출력
        if(a==1):
            update(0,n-1,1,b-1,c)
            nodes[b-1]=c
            # print(tree,'\n',nodes)
        else:
            print(query(0,n-1,1,b-1,c-1))