[BOJ] 2042 - 구간 합 구하기

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

BOJ - 구간 합 구하기

문제 유형: 세그먼트 트리

문제 난이도: Gold I

시간 제한: 2초

메모리 제한: 256MB

 

문제

어떤 N개의 수가 주어져 있다. 그런데 중간에 수의 변경이 빈번히 일어나고 그 중간에 어떤 부분의 합을 구하려 한다. 만약에 1,2,3,4,5 라는 수가 있고, 3번째 수를 6으로 바꾸고 2번째부터 5번째까지 합을 구하라고 한다면 17을 출력하면 되는 것이다. 그리고 그 상태에서 다섯 번째 수를 2로 바꾸고 3번째부터 5번째까지 합을 구하라고 한다면 12가 될 것이다.

 

입력

첫째 줄에 수의 개수 N(1 ≤ N ≤ 1,000,000)과 M(1 ≤ M ≤ 10,000), K(1 ≤ K ≤ 10,000) 가 주어진다. M은 수의 변경이 일어나는 횟수이고, K는 구간의 합을 구하는 횟수이다. 그리고 둘째 줄부터 N+1번째 줄까지 N개의 수가 주어진다. 그리고 N+2번째 줄부터 N+M+K+1번째 줄까지 세 개의 정수 a, b, c가 주어지는데, a가 1인 경우 b(1 ≤ b ≤ N)번째 수를 c로 바꾸고 a가 2인 경우에는 b(1 ≤ b ≤ N)번째 수부터 c(b ≤ c ≤ N)번째 수까지의 합을 구하여 출력하면 된다.

입력으로 주어지는 모든 수는 -2^63보다 크거나 같고, 2^63-1보다 작거나 같은 정수이다.

 

출력

첫째 줄부터 K줄에 걸쳐 구한 구간의 합을 출력한다. 단, 정답은 -2^63보다 크거나 같고, 2^63-1보다 작거나 같은 정수이다.

 

풀이

세그먼트 트리의 입문 문제이다.

https://www.youtube.com/watch?v=-dUiRtJ8ot0

배열로 이진트리를 표현해준다.

원본 배열이 n의 크기를 가진다면,

세그먼트 트리는 최대 4*n의 크기를 가진다.

 

이진트리의 배열 표현상, 왼쪽자식은 인덱스*2, 오른쪽 자식은 인덱스*2+1이다. (1-indexed 기준, 0-indexed시 1을 더 더해주면 된다.)

seg[index]에는 [l, r]구간의 연산이 저장되어 있다. 이 문제에서는 양쪽 자식들의 합이다.

 

build:

l == r인경우, 리프 노드이다.

seg[index] = arr[left]이다. 대입 후 리턴한다.

 

아닌 경우, mid를 구하여서, 양쪽 자식 서브트리들을 먼저 생성한다.

왼쪽자식의 구간은 [l, mid], 오른쪽자식의 구간은 [mid+1, r]이다.

두 자식들의 값이 구해지면, 두 자식들의 합을 seg[index]에 저장한다.

 

update:

build와 유사하다.

원본 배열을 바꾸고 새로 빌드한다고 보면 된다.

 

query(여기서는 subarrSum이라고 되어있다.):

query에는 쿼리의 범위를 받아야 한다, ql, qr이라 하겠다.

우선, ql이 r보다 크거나 qr이 l보다 작으면 안된다. 여기서는 구간합이므로 0을 반환해야 영향을 끼치지 않는 것이 되겠다.

만약 ql이 l보다 작거나 같고, qr이 r보다 크거나 같으면, 현재 범위 [l, r]이 쿼리 범위 안에 드는 가장큰 구간이므로, 여기서 반환한다.

 

그게 아니라면, 분할해야 한다.

mid를 구해서, 왼쪽부분에 대해, 오른쪽부분에 대해 쿼리를 날리고, 돌아온 두 응답을 합쳐서 리턴한다.

 

코드

C++

#include <bits/stdc++.h>
#define LIMIT 1000001
#define ll long long
using namespace std;

int n, m, k;

ll a[LIMIT], seg[4 * LIMIT];

void build(int index, int left, int right) {
  if (left == right) {
    seg[index] = a[left];
    return;
  }

  int mid = (left + right) / 2;
  build(index * 2, left, mid);
  build(index * 2 + 1, mid + 1, right);

  seg[index] = seg[index * 2] + seg[index * 2 + 1];
}

void update(int b, ll value, int index, int left, int right) {
  if (left == right) {
    seg[index] = value;
    return;
  }
  int mid = (left + right) / 2;
  if (b <= mid) {
    update(b, value, index * 2, left, mid);
  } else {
    update(b, value, index * 2 + 1, mid + 1, right);
  }
  seg[index] = seg[index * 2] + seg[index * 2 + 1];
}

ll subarrSum(int index, int qleft, int qright, int left, int right) {
  if (right < qleft || left > qright) {
    return 0;
  }
  if (left >= qleft && right <= qright) {
    return seg[index];
  }
  int mid = (left + right) / 2;
  ll resLeft = subarrSum(index * 2, qleft, qright, left, mid);
  ll resRight = subarrSum(index * 2 + 1, qleft, qright, mid + 1, right);

  return resLeft + resRight;
}

int main(int argc, char *argv[]) {
  ios_base::sync_with_stdio(0);
  cin.tie(0);

  cin >> n >> m >> k;
  for (int i = 1; i <= n; ++i) {
    cin >> a[i];
  }

  build(1, 1, n);

  for (int i = 0; i < m + k; ++i) {
    int x, y;
    ll z;
    cin >> x >> y >> z;
    if (x == 1) {
      update(y, z, 1, 1, n);
    } else {
      cout << subarrSum(1, y, z, 1, n) << "\n";
    }
  }

  return 0;
}

 

Go

package main

import (
	"bufio"
	"fmt"
	"os"
)

var (
	n, m, k int32
	arr     []int64
	seg     []int64
)

func build(index, l, r int32) {
	if l == r {
		seg[index] = arr[l]
		return
	}

	mid := (l + r) / 2
	build(index*2, l, mid)
	build(index*2+1, mid+1, r)

	seg[index] = seg[index*2] + seg[index*2+1]
}

func update(index, target int32, newVal int64, l, r int32) {
	if l == r {
		arr[target] = newVal
		seg[index] = newVal
		return
	}

	mid := (l + r) / 2
	if target <= mid {
		update(index*2, target, newVal, l, mid)
	} else {
		update(index*2+1, target, newVal, mid+1, r)
	}
	seg[index] = seg[index*2] + seg[index*2+1]
}

func subarrSum(index, ql, qr, l, r int32) int64 {
	if qr < l || ql > r {
		return 0
	}

	if ql <= l && r <= qr {
		return seg[index]
	}

	mid := (l + r) / 2
	lRes := subarrSum(index*2, ql, qr, l, mid)
	rRes := subarrSum(index*2+1, ql, qr, mid+1, r)

	return lRes + rRes
}

func main() {
	reader := bufio.NewReader(os.Stdin)
	writer := bufio.NewWriter(os.Stdout)
	defer writer.Flush()

	fmt.Fscan(reader, &n, &m, &k)
	arr = make([]int64, n+1)
	seg = make([]int64, 4*(n+1))

	for i := int32(1); i <= n; i++ {
		fmt.Fscan(reader, &arr[i])
	}

	build(1, 1, n)

	for i := int32(0); i < m+k; i++ {
		var a, b int32
		var c int64
		fmt.Fscan(reader, &a, &b, &c)

		if a == 1 {
			update(1, b, c, 1, n)
		} else {
			res := subarrSum(1, b, int32(c), 1, n)
			fmt.Fprintln(writer, res)
		}
	}

}