[BOJ] 10999 - 구간 합 구하기 2

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

BOJ - 구간 합 구하기 2

문제 유형: 세그먼트 트리 (with lazy propagation)

문제 난이도: Platinum IV

시간 제한: 2초

메모리 제한: 256MB

 

문제

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

 

입력

첫째 줄에 수의 개수 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, b, c, d가 주어지는데, a가 1인 경우 b번째 수부터 c번째 수에 d를 더하고, a가 2인 경우에는 b번째 수부터 c번째 수의 합을 구하여 출력하면 된다.

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

 

출력

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

 

풀이

세그먼트 트리 with lazy propagation의 정석인 문제이다.

세그먼트 트리, lazy propagation이란?

https://riveroverflow.tistory.com/entry/%EC%84%B8%EA%B7%B8%EB%A8%BC%ED%8A%B8-%ED%8A%B8%EB%A6%AC%EC%99%80-%EC%A7%80%EC%97%B0%EC%A0%84%ED%8C%8C-Segment-Tree-with-lazy-propagation

세그먼트 트리와 지연전파 (Segment Tree, with lazy propagation)

 

정수 범위와 인덱스에 유의해서 구현해주면 된다.

 

코드

Go

package main

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

var (
	arr  []int64
	seg  []int64
	lazy []int64
)

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

	mid := (l + r) >> 1
	build(index*2+1, l, mid)
	build(index*2+2, mid+1, r)
	seg[index] = seg[index*2+1] + seg[index*2+2]
}

func propagate(index, l, r int32) {
	if lazy[index] != 0 {
		seg[index] += lazy[index] * int64(r-l+1)
		if l != r {
			lazy[index*2+1] += lazy[index]
			lazy[index*2+2] += lazy[index]
		}
		lazy[index] = 0
	}
}

func update(index, l, r, ql, qr int32, val int64) {
	propagate(index, l, r)
	if l > qr || r < ql {
		return
	}

	if ql <= l && r <= qr {
		lazy[index] += val
		propagate(index, l, r)
		return
	}

	mid := (l + r) >> 1
	update(index*2+1, l, mid, ql, qr, val)
	update(index*2+2, mid+1, r, ql, qr, val)
	seg[index] = seg[index*2+1] + seg[index*2+2]
}

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

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

	mid := (l + r) >> 1
	lres := query(index*2+1, l, mid, ql, qr)
	rres := query(index*2+2, mid+1, r, ql, qr)
	return lres + rres
}

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

	var n, m, k int32
	fmt.Fscan(reader, &n, &m, &k)

	arr = make([]int64, n)
	seg = make([]int64, n*4)
	lazy = make([]int64, n*4)

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

	build(0, 0, n-1)
	for i := int32(0); i < m+k; i++ {
		var a, b, c int32
		var d int64
		fmt.Fscan(reader, &a)
		if a == 1 {
			fmt.Fscan(reader, &b, &c, &d)
			b--
			c--
			update(0, 0, n-1, b, c, d)
		} else {
			fmt.Fscan(reader, &b, &c)
			b--
			c--
			fmt.Fprintln(writer, query(0, 0, n-1, b, c))
		}
	}

}