본문 바로가기

컴퓨터/코테

백준 2042 세그먼트 트리

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

nums = [0] + [int(sys.stdin.readline()) for _ in range(N)]
position = [0 for _ in range(N+1)]
orders = [list(map(int, sys.stdin.readline().split())) for _ in range(M+K)]
arr = [0] * 4 * N

def init(start = 1, end = N, index = 1):
    if start == end :
        arr[index] = nums[start]
        position[start] = index
        return arr[index]
    mid = (start + end) // 2
    arr[index] = init(start, mid, index*2) + init(mid+1, end, index*2 + 1)
    return arr[index]

def getSum(Left, Right, start=1, end=N, index=1):
    if Left <= start and end <= Right :
        return arr[index]
    if start > Right or end < Left :
        return 0
    mid = (start + end)//2
    return getSum(Left, Right, start, mid, index*2) + getSum(Left, Right, mid+1, end, index*2 + 1)

def updateArr(idx):
    arr[idx] = arr[idx*2] + arr[idx*2 + 1]
    if idx == 1 :
        return
    updateArr(idx//2)

init()
for typeOfOrder, A, B in orders :
    if typeOfOrder == 2:
        #구간합
        print(getSum(A,B))
    else :
        #업데이트
        arr[position[A]] = B
        updateArr(position[A]//2)

오랜만에 뭔가 재밌게 풀었다고 할만한 문제였다. 세그먼트 트리 관련 참고자료가 많아서 공부도 수월하게 하면서 문제 풀 수 있었다. 분명 모든 연산이 logN 인데도 시간초과가 떠서 뭐지 싶었는데 알고보니 input()이 원인이었다. sys.stdin.readline()을 써서 해결했다.

'컴퓨터 > 코테' 카테고리의 다른 글

릿코드 37- Sudoku Solver, nQueen 다시풀기  (0) 2021.09.01
nQueen - 백트래킹  (0) 2021.09.01
백준 17114 - 미세먼지 확산  (0) 2021.08.30
백준 11657 타임머신  (0) 2021.08.30
카카오 기출 - 카드 뒤집기  (0) 2021.08.27