백준 하늘에서 떨어지는 1, 2, ..., R-L+1개의 별(17353)

하늘에서 떨어지는 1, 2, ..., R-L+1개의 별

1. 힌트

1) 구간 합 세그먼트 트리를 구성하고 $1$번 구간 업데이트 쿼리를 lazy propagation으로 처리한다.

2) $2$번 쿼리가 구간 쿼리가 아닌 점 쿼리임에 유의한다.

3) 점 쿼리이므로 구간 $S$의 정의를 구간의 이전 원소와의 차이로 정의하면 $S[i]$의 값은 $pSum[i]$로 구할 수 있다.

2. 접근

1) 비슷한 문제를 풀어본 적이 있던가?

수열과 쿼리 21

점 쿼리만 들어온다면 주어진 수열을 이전 수열과의 차이만을 저장하게 바꿔도 점 쿼리에 대답할 수 있다.
구간 합 세그먼트 트리를 이렇게 구현한다면 $1$번 쿼리를 lazy propagation을 이용해 간단하게 구현할 수 있다.

3. 구현

import java.util.*;
import java.io.*;

public class Main {
	
	public static void main(String[] args) throws IOException {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		StringBuilder bw = new StringBuilder();
		int N = Integer.parseInt(br.readLine());
		int[] S = new int[N];
		StringTokenizer st = new StringTokenizer(br.readLine(), " ");
		for (int i = 0; i < N; i++) S[i] = Integer.parseInt(st.nextToken());
		int[] A = new int[N]; // 이전 원소와의 차이 수열
		A[0] = S[0];
		for (int i = 1; i < N; i++) A[i] = S[i] - S[i - 1];
		RSQ q = new RSQ(A);
		int Q = Integer.parseInt(br.readLine());
		for (int i = 0; i < Q; i++) {
			st = new StringTokenizer(br.readLine(), " ");
			int o = Integer.parseInt(st.nextToken());
			if (o == 1) {
				int L = Integer.parseInt(st.nextToken()) - 1;
				int R = Integer.parseInt(st.nextToken()) - 1;
				q.update(L, R, 1);
				q.update(R + 1, R + 1, -(R - L + 1));
			} else {
				bw.append(q.query(0, Integer.parseInt(st.nextToken()) - 1)).append("\n");
			}
		}
		System.out.print(bw);
	}
	
}
class RSQ {
	int n;
	long[] range; long[] lazy;
	
	RSQ(int[] array) {
		n = array.length;
		range = new long[4 * n];
		lazy = new long[4 * n];
		init(array, 1, 0, n - 1);
	}
	
	long init(int[] array, int node, int nodeLeft, int nodeRight) {
		if (nodeLeft == nodeRight) return range[node] = array[nodeLeft];
		int mid = (nodeLeft + nodeRight) / 2;
		return range[node] = init(array, 2 * node, nodeLeft, mid) + init(array, 2 * node + 1, mid + 1, nodeRight);
	}
	
	long query(int left, int right) {
		return query(left, right, 1, 0, n - 1);
	}
	
	long query(int left, int right, int node, int nodeLeft, int nodeRight) {
		propagate(node, nodeLeft, nodeRight);
		if (right < nodeLeft || nodeRight < left) return 0;
		if (left <= nodeLeft && nodeRight <= right) return range[node];
		int mid = (nodeLeft + nodeRight) / 2;
		return query(left, right, 2 * node, nodeLeft, mid) + query(left, right, 2 * node + 1, mid + 1, nodeRight);
	}
	
	void update(int left, int right, int k) {
		update(left, right, k, 1, 0, n - 1);
	}
	
	void update(int left, int right, int k, int node, int nodeLeft, int nodeRight) {
		propagate(node, nodeLeft, nodeRight);
		if (right < nodeLeft || nodeRight < left) return;
		if (left <= nodeLeft && nodeRight <= right) {
			lazy[node] += k;
			propagate(node, nodeLeft, nodeRight);
			return;
		}
		int mid = (nodeLeft + nodeRight) / 2;
		update(left, right, k, 2 * node, nodeLeft, mid);
		update(left, right, k, 2 * node + 1, mid + 1, nodeRight);
		range[node] = range[2 * node] + range[2 * node + 1];
	}
	
	void propagate(int node, int nodeLeft, int nodeRight) {
		if (lazy[node] != 0) {
			if (nodeLeft != nodeRight) {
				lazy[2 * node] += lazy[node];
				lazy[2 * node + 1] += lazy[node];
			}
			range[node] += lazy[node] * (nodeRight - nodeLeft + 1);
			lazy[node] = 0;
		}
	}
	
}

1) 시간 복잡도

두 가지 쿼리 모두 시간 복잡도가 $O(lgN)$이므로 전체 시간 복잡도는 $O(QlgN)$

2) 테스트 케이스