백준 나무 심기(1280)

나무 심기

오늘은 식목일이므로 나무 심기 문제에 대해서 풀이해본다.
나무를 심는데 드는 비용은 현재 심어져있는 모든 나무 까지의 거리의 합이라고 한다. 그런데 나무는 어차피 문제에 주어진 순서대로 심어야 한다니깐 심는 순서는 고려할 필요가 없고,
특이하게도 2번 나무부터 $N$번 나무까지 심는 비용의 곱을 출력하라는데 $i$번째 나무를 심어야 할 때 $i - 1$번째 나무를 심는데 드는 비용을 알면 좀 더 간단하게 할 수 있을까?

1. 그림으로 그려볼 수 있을까?

$\{1, 2, 3, 4, 6, 7, 8\}$위치에 나무가 심어져있다고 할 때, $X[i + 1]$이 6이라면 $cost(i)$와 $cost(i + 1)$이 동일해지긴 하는데 일반적인 상황에서 계속 성립하는 것은 아닌것 같다.

2. 내가 문제를 푸는 과정을 수식화할 수 있을까?

$i~(i \ge2)$번째 나무를 심는데 드는 비용은 다음과 같이 나타낼 수 있다. (1-based)
$cost(i) =\displaystyle\sum_{k = 1}^{i - 1}|X[i]~-~X[k]|~~(i \ge1)$

이 값을 계산하기 위해서는 $O(N)$으로 이전까지 심은 나무들을 모두 봐야하는데 여기서 $X[i]$는 $k$와 상관이 없으므로 $\Sigma$밖으로 나올 수 있다는 점을 깨달았다. 밖으로 빼내기만 한다면 각각은 세그먼트 트리를 이용해서 $O(lgN)$으로 구할 수 있다.

3. 문제를 분해할 수 있을까?

$cost(i) =\displaystyle\sum_{k = 1}^{i - 1}|X[i]~-~X[k]|~~(i \ge1)$
이 수식에서 $X[i]$를 빼내려면 절댓값 기호를 없애야하는데 이는 범위를 나눠서 해결할 수 있다.
$[1, i - 1]$에 해당하는 $k$의 범위를
$X[k] \le X[i]$인 경우와 $X[k] \gt X[i]$로 나누어보자
$cost(i) =\displaystyle\sum_{k = 1}^{i - 1}X[i]~-~X[k]~~(i \ge1, X[k] \le X[i])$
$~~~~~~~~~~~~~+\displaystyle\sum_{k = 1}^{i - 1}X[k]~-~X[i]~~(i \ge1,X[k] \gt X[i])$

이렇게 되면 각각의 $\Sigma$에서 $X[i]$를 빼낼 수 있는데 이는 $X[i]$의 왼쪽에(같은 경우도 포함) 존재하는 나무의 수에다가 $X[i]$를 곱한 값이 된다.
$X[i]$를 빼내고 나면 남은 $X[k]$는 $\displaystyle-\sum_{k = 1}^{i - 1}X[k]~~(i \ge1)$ 이 되므로 이는 $X[i]$의 왼쪽에(같은 경우도 포함) 존재하는 나무의 좌표값들의 합과 같다.

4. 세그먼트 트리

전체 구간을 두 구간으로 나누어서($[0, X[i]]$와 $[X[i] + 1, 200000]$) 구간에 대해서 쿼리가 들어오고, 쿼리의 업데이트를 $O(N)$보다 빠르게 처리하려면 세그먼트 트리가 필요하다.
그런데 필요한 쿼리는 구간 내에 좌표의 개수와 좌표값들의 합이므로 구간합을 간단하게 처리할 수 있는 펜윅 트리를 이용해서 해결한다.

5. 구현

import java.io.*;

public class Main {
	static final int MOD = 1000000007;
	
	public static void main(String[] args) throws IOException {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		int N = Integer.parseInt(br.readLine());
		// Range Count Query [0, pos]에 존재하는 좌표의 수를 답한다
		FenwickTree rcq = new FenwickTree(200001); 
		// Range Sum Query [0, pos]에 존재하는 좌표값들의 합을 답한다
		FenwickTree rsq = new FenwickTree(200001);
		int ret = 1;
		int x = Integer.parseInt(br.readLine());
		rcq.add(x, 1); rsq.add(x, x);
		for (int i = 2; i <= N; i++) {
			x = Integer.parseInt(br.readLine());
            		// X[i]를 기준으로 두 구간으로 나눈다.
			// [0, x]
			long temp = (rcq.sum(x) * x - rsq.sum(x) + MOD) % MOD;
			// [x, 200000]
			temp = (temp + (rsq.sum(200000) - rsq.sum(x)) - (rcq.sum(200000) - rcq.sum(x)) * x + MOD) % MOD;
			rcq.add(x, 1); rsq.add(x, x);
			ret = (int)((ret * temp) % MOD);
		}
		System.out.println(ret);
	}
	
}

//펜윅 트리
class FenwickTree {
	long[] tree;
	public FenwickTree(int n) {tree = new long[n + 1];}

	void add(int pos, int val) {
		// 내부적으로는 1-based
		pos++;
		while (pos < tree.length) {
			tree[pos] += val;
			pos += (pos & -pos);
		}
	}

	long sum(int pos) {
		// 내부적으로는 1-based
		pos++;
		long ret = 0;
		while (pos > 0) {
			ret += tree[pos];
			pos &= (pos - 1);
		}
		return ret;
	}

}