본문 바로가기
Problem Solving/Segment Tree

[세그먼트 트리/파이썬] 백준 2042번 : 구간 합 구하기 / 골드 1

by ggyongi 2022. 3. 27.
반응형

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

 

2042번: 구간 합 구하기

첫째 줄에 수의 개수 N(1 ≤ N ≤ 1,000,000)과 M(1 ≤ M ≤ 10,000), K(1 ≤ K ≤ 10,000) 가 주어진다. M은 수의 변경이 일어나는 횟수이고, K는 구간의 합을 구하는 횟수이다. 그리고 둘째 줄부터 N+1번째 줄

www.acmicpc.net

import sys
input = sys.stdin.readline
n, m, k = map(int, input().split())
lst = []
for _ in range(n):
    lst.append(int(input()))

def merge(left, right):
    return left + right

def build(stree, node, left, right):
    # leaf node
    if left == right:
        stree[node] = lst[left]
        return stree[node]

    mid = left + (right - left)//2
    left_val = build(stree, 2*node, left, mid)
    right_val = build(stree, 2*node + 1, mid + 1, right)
    stree[node] = merge(left_val, right_val)
    return stree[node]

def query(start, end, node, left, right):
    # 나의 담당구역 left~right 사이에 start~end 가 아예 포함되지 않는 경우 => 가지 치기
    if end < left or start > right:
        return 0

    # 나의 담당구역 left~right 전체가 start~end 에 포함되는 경우 => 내 정보를 바로 리턴
    if start <= left and right <= end:
        return stree[node]

    mid = left + (right - left) // 2
    left_val = query(start, end, 2 * node, left, mid)
    right_val = query(start, end, 2 * node + 1, mid + 1, right)
    return merge(left_val, right_val)

def update(idx, val, node, left, right):
    # 나의 담당구역 left~right 사이에 target이 없는 경우 => 기존 값을 그대로 반환(변경 필요 x)
    if idx < left or idx > right:
        return stree[node]

    # leaf node: 바꿀 idx 값 발견 => 값 변경
    if left == right:
        stree[node] = val
        return stree[node]

    mid = left + (right - left) // 2
    left_val = update(idx, val, 2*node, left, mid)
    right_val = update(idx, val, 2*node + 1, mid + 1, right)
    stree[node] = merge(left_val, right_val)
    return stree[node]

stree = [0 for _ in range(4*n)]
build(stree, 1, 0, n-1)

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

세그먼트 트리 자료구조를 구현하면 쉽게 풀린다.

 

비전공자 네카라 신입 취업 노하우

시행착오 끝에 얻어낸 취업 노하우가 모두 담긴 전자책!

kmong.com

댓글