[BOJ] 16975 - 수열과 쿼리 21

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

BOJ - 수열과 쿼리 21

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

문제 난이도: Platinum IV

시간 제한: 2초

메모리 제한: 512MB

 

문제

길이가 N인 수열 A1, A2, ..., AN이 주어진다. 이때, 다음 쿼리를 수행하는 프로그램을 작성하시오.

• 1 i j k: Ai, Ai+1, ..., Aj에 k를 더한다.
• 2 x: Ax 를 출력한다.

 

입력

첫째 줄에 수열의 크기 N (1 ≤ N ≤ 100,000)이 주어진다.

둘째 줄에는 A1, A2, ..., AN이 주어진다. (1 ≤ Ai ≤ 1,000,000)

셋째 줄에는 쿼리의 개수 M (1 ≤ M ≤ 100,000)이 주어진다.

넷째 줄부터 M개의 줄에는 쿼리가 한 줄에 하나씩 주어진다. 1번 쿼리의 경우 1 ≤ i ≤ j ≤ N, -1,000,000 ≤ k ≤ 1,000,000 이고, 2번 쿼리의 경우 1 ≤ x ≤ N이다. 2번 쿼리는 하나 이상 주어진다.

 

출력

2번 쿼리가 주어질 때마다 출력한다.

 

풀이

lazy propagation을 이용하여 구간업데이트 쿼리를 최적화하고, 조회 쿼리에 대해서는 이분탐색 like로 풀 수 있다.

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 (
	a, seg, lazy []int64
)

func build(index, l, r int32) {
	if l == r {
		seg[index] = a[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 rangeUpdate(index, l, r, ql, qr, inc int32) {
	propagate(index, l, r)
	if ql > r || qr < l {
		return
	}

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

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

func query(index, l, r, target int32) int64 {
	propagate(index, l, r)
	if l == r {
		return seg[index]
	}
	mid := (l + r) >> 1
	if target <= mid {
		return query(index*2+1, l, mid, target)
	} else {
		return query(index*2+2, mid+1, r, target)
	}
}

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

	var n, m int32
	fmt.Fscan(reader, &n)
	a = make([]int64, n)
	seg = make([]int64, 4*n)
	lazy = make([]int64, 4*n)

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

	build(0, 0, n-1)

	for i := int32(0); i < m; i++ {
		var cmd int32
		fmt.Fscan(reader, &cmd)
		if cmd == 1 {
			var ql, qr, k int32
			fmt.Fscan(reader, &ql, &qr, &k)
			rangeUpdate(0, 0, n-1, ql-1, qr-1, k)
		} else {
			var x int32
			fmt.Fscan(reader, &x)
			fmt.Fprintln(writer, query(0, 0, n-1, x-1))
		}
	}
}