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

 

[문제]

 

www.acmicpc.net/problem/17353

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

 

 

[난이도]

 

- Platinum II (solved.ac 20.12.22 기준)

 

 

[필요 개념]

 

 - Lazy Propagation with Segment Tree

 - Fenwick Tree

 

 

디스크립션은 간단하다. 

두 종류의 쿼리가 들어오고, 한 쿼리는 구간 [L, R]에 순서대로 1, 2, ..., R-L+1을 더해준다.

나머지 한 쿼리는 X번째 원소의 값을 출력해준다.

 

 

[풀이 1]

 

구간 업데이트를 해주는 것은 lazy propagation이지만, 구간의 원소들에 모두 같은 값을 더해주는 것이 아니다.

단순하게 레이지 세그를 이용해서 각 원소에 모두 다른 값을 더해주기는 어렵기 때문에 조금 변형시켜 

구간에 모두 같은 값을 더해주는 형태로 바꿔주면 레이지 세그를 이용할 수 있을 것이다. 

 

쿼리 1을 보면 구간에 더해지는 값들은 공차가 1인 등차수열을 이룬다. 

따라서 이웃한 항과의 차가 항상 1이므로 원래의 수열을 A[i]라고 하면, B[i] = A[i] - A[i-1] 수열을 생각해본다. 

이렇게 되면 [L, R] 구간에 모두 1을 더해주는 것과 쿼리 1이 동일한 의미를 가진다. 

B[L+1] + 1 == A[L+1] - A[L] + 1 == A[L+1] + 2 - (A[L] + 1) 를 통해서 이해해보자.

 

이 때, 주의해야 할 부분은 실제로 A[R]에는 R-L+1만큼 더해주었기 때문에 B[R+1] = A[R+1] - A[R]은 R-L+1만큼 줄어들게 된다. 

그러므로 B[R+1]은 R-L+1만큼 감소시켜주어야 한다. 

 

쿼리 2는 A[X]를 구해야 하는데, B[X] + B[X-1] + ... + B[1] = (A[X] - A[X-1]) + (A[X-1] - A[X-2]) + ... + (A[1] - A[0]) 

= A[X] - A[0] = A[X] 를 만족하기 때문에 [1, X] 구간의 수열 B의 원소 합을 구해주면 된다.

 

 

[소스 코드]

 

#include<bits/stdc++.h>
using namespace std;
using ll = long long;
ll A[100001];
ll tr[400001];
ll lazy[400001];
void update_lazy(int x, int s, int e);
ll sum(int x, int s, int e, int l, int r) {
	update_lazy(x, s, e);
	if (s > r || e < l) return 0;
	else if (s >= l && e <= r) return tr[x];
	else return sum(x * 2, s, (s + e) / 2, l, r) + sum(x * 2 + 1, (s + e) / 2 + 1, e, l, r);
}
void update_range(int x, int s, int e, int l, int r, ll val) {
	update_lazy(x, s, e);
	if (s > r || e < l) return;
	if (s >= l && e <= r) {
		tr[x] += (e - s + 1)*val;
		if (s != e) {
			lazy[x * 2]+= val;
			lazy[x * 2 + 1]+= val;
		}
		return;
	}
	update_range(x * 2, s, (s + e) / 2, l, r, val);
	update_range(x * 2 + 1, (s + e) / 2 + 1, e, l, r, val);
	tr[x] = tr[x * 2] + tr[x * 2 + 1];
}
void update_lazy(int x, int s, int e) {
	if (lazy[x] != 0) {
		tr[x] += (e - s + 1)*lazy[x];
		if (s != e) {
			lazy[x * 2] += lazy[x];
			lazy[x * 2 + 1] += lazy[x];
		}
		lazy[x] = 0;
	}
}
int main(void) {
	ios::sync_with_stdio(false); cin.tie(nullptr); cout.tie(nullptr);
	int N; cin >> N;
	for (int i = 1; i <= N; i++) cin >> A[i];
	for (int i = 1; i <= N; i++) {
		update_range(1, 1, N, i, i, A[i] - A[i-1]);
	}
	int Q; cin >> Q;
	while (Q--) {
		int a; cin >> a;
		if (a == 1) {
			int l, r; cin >> l >> r;
			update_range(1, 1, N, l, r, 1);
			update_range(1, 1, N, r + 1, r + 1, -(r - l + 1));
		}
		else {
			int x; cin >> x;
			cout << sum(1, 1, N, 1, x) << '\n';
		}
	}
}

 

[풀이 2]

 

Range update & Point query (구간 업데이트 & 점 쿼리)면 펜윅 트리를 사용할 수 있다. 

다만, 마찬가지로 구간에 모두 같은 값을 더해주는 것이 아니기 때문에 추가적인 작업이 필요하다. 

 

구간 [L, R]에 1, 2, ..., R-L+1을 더해주기 때문에, [L, R] 사이의 인덱스 k에 대해서는 k - L + 1을 더해주는 것과 동일하다. 

그렇다면, A[k]에는 이 원소를 포함하는 구간 쿼리들에 대해서 L값만 고려를 해줘도 충분하다. 

예를 들어서 k = 4라고 하고, [3, 6] 구간 업데이트를 하면 A[k]에는 4 - 3 + 1이 더해지게 된다.

그리고 [1, 5] 구간 업데이트를 하면 4 - 1 + 1이 더해지게 된다.

이후에 A[k]의 값을 구한다면 A[k] + (4 - 3 + 1) + (4 - 1 + 1) = A[k] + 6이 나오게 된다.

즉, 최종적으로 A[k]를 구하게 되면 'A[k]의 update 수 * (k+1) - L값들의 합' 이 더해진 값으로 나온다. 

 

따라서, 각 원소별로 업데이트된 횟수(cnt)와 L값들의 합(tr)을 각각 저장하는 배열을 만들어서 관리해주면 된다.

 

[소스 코드]

 

#include<bits/stdc++.h>
using namespace std;
using ll = long long;
using pll = pair<ll, ll>;
ll A[100001];
ll tr[100001];
ll cnt[100001];
int N, Q;
void update(int i, ll val) {
	while (i <= N) {
		tr[i] += val;
		if(val > 0) cnt[i]++;
		else cnt[i]--;
		i += (i & -i);
	}
}
pll sum(int i) {
	pll ans = { 0, 0 };
	while (i > 0) {
		ans.first += tr[i];
		ans.second += cnt[i];
		i -= (i & -i);
	}
	return ans;
}
int main(void) {
	ios::sync_with_stdio(false); cin.tie(nullptr); cout.tie(nullptr);
	cin >> N;
	for (int i = 1; i <= N; i++) cin >> A[i];
	cin >> Q;
	while (Q--) {
		int a; cin >> a;
		if (a == 1) {
			int l, r; cin >> l >> r;
			update(l, l);
			update(r + 1, -l);
		}
		else {
			int x; cin >> x;
			pll p = sum(x);
			cout << p.second * (x + 1) - p.first + A[x] << '\n';
		}
	}
}

 

PC로 보시는 것을 권장합니다. 

피드백은 언제나 환영입니다. 댓글로 달아주세요 ^-^