[백준/2268] 수들의 합 7 (Python)

Problem : https://www.acmicpc.net/problem/2268

2268번: 수들의 합 7

 

Difficulty : Gold 1

 

Status : Solved

 

Time : 00:21:48

 

 

---

 

문제 설명

 

 

N개의 수 A[1], A[2], …, A[N] 이 주어졌을 때, 함수 Sum(i, j)는 A[i] + A[i+1] + … + A[j]를 구하는 함수이다. (i > j일 경우에는 A[j] + A[j+1] + ... + A[i]) A가 주어졌을 때, Sum(i, j)를 구하는 것은 매우 쉬운 문제이다. 이러한 (i, j)가 여러 개 주어졌을 때도 별로 어려운 문제는 아니다.

Sum함수와 더불어 Modify라는 함수를 정의하자. Modify(i, k)를 수행하면 A[i] = k가 되는 함수이다. Sum함수와 Modify 함수의 사용 목록이 주어졌을 때, 이에 해당하는 연산을 하는 프로그램을 작성하시오. 두 함수를 섞어서 사용할 수도 있다.

 

 

입력 및 출력

 

입력

 

첫째 줄에는 N(1 ≤ N ≤ 1,000,000), M(1 ≤ M ≤ 1,000,000)이 주어진다. M은 수행한 명령의 개수이며 다음 M개의 줄에는 수행한 순서대로 함수의 목록이 주어진다. 첫 번째 숫자는 어느 함수를 사용했는지를 나타내며, 0일 경우에는 Sum 함수를, 1일 경우에는 Modify 함수를 나타낸다. 다음 두 수는 각 함수의 인자 (i, j)나 (i, k)를 나타낸다. 처음에는 A[1] = A[2] = … = A[N] = 0이다. Modify인 경우에 1 ≤ k ≤ 100,000 이다.

 

출력

 

Sum 함수의 개수만큼 각 줄에 Sum 함수의 리턴값을 출력한다.

 

입력 예시

 

3 5
0 1 3
1 1 2
1 2 3
0 2 3
0 1 3

 

출력 예시

 

0
3
5

 

 

---

 

풀이

 

부분합을 구하는 데 O(logN), 업데이트하는데 O(logN)이 필요하다면, 일단은 세그먼트 트리를 생각해내보자. 잊을 만하면 튀어나오는 자료구조라 복습 겸 풀어보았다 :)

 

 

풀이 코드

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

class SegTree :
  def __init__(self) :
    self.tree = [0]*(N*4+1)

  def modify(self, i, k) :
    self._modify(0, N-1, 1, i-1, k)
    
  def _modify(self, start, end, idx, target, val) :
    if target < start or target > end :
      return
    if start == end == target :
      self.tree[idx] = val
      return

    mid = (start + end) // 2
    self._modify(start, mid, 2*idx, target, val)
    self._modify(mid+1, end, 2*idx+1, target, val)
    self.tree[idx] = sum(self.tree[2*idx:2*idx+2])

  def sum(self, i, j) :
    return self._sum(0, N-1, 1, i-1, j-1)

  def _sum(self, start, end, idx, l, r) :
    if l <= start <= end <= r :
      return self.tree[idx]
    if r < start or end < l :
      return 0

    mid = (start + end) // 2
    lval = self._sum(start, mid, 2*idx, l, r)
    rval = self._sum(mid+1, end, 2*idx+1, l, r)
    return lval + rval

segtree = SegTree()
for _ in range(M) :
  q, *cmd = map(int, input().split())
  if q == 0 :
    i, j = cmd
    if i > j :
      i, j = j, i
    print(segtree.sum(i, j))
  else :
    i, k = cmd
    segtree.modify(i, k)

풀이 완료!