[P4] 백준 1280번 나무 심기 C++ 세그먼트 트리

리뷰

 

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

세그먼트 트리를 활용하여 쿼리 순으로 배열 내 모든 요소와의 거리를 구하는 문제

 

 

전역 변수

• N : 좌표의 최대값을 저장할 정수형 변수
• MOD : 모듈러 값을 지정할 정수형 상수 변수
• n : 나무의 개수를 저장할 정수형 변수
• tree : 세그먼트 트리를 통해 구간 누적합과 누적 개수를 저장할 pair<long long, int>타입의 배열

 

함수

1. update

void update(int node, int s, int e, ll idx)

 

나무를 심고 세그먼트 트리 정보를 업데이트 하는 문제

1. 매개 변수로 노드 정보 node, 탐색 범위 s, e, 나무의 좌표 idx를 전달받는다.
2. 기저 조건으로 리프 노드에 도달했을 경우 현재 노드에 나무의 위치와 개수를 증가시킨다.
3. 리프 노드가 아닐 경우 idx가 위치한 쪽으로 재귀 탐색을 진행해 준다.
4. 재귀를 빠져나오며 상위 노드를 업데이트 해준다.

 

2. query

pair<ll, int> query(int node, int s, int e, int L, int R)

 

심은 나무를 기준으로 좌, 우를 탐색하여 누적합과 개수를 구하는 함수

1. 매개 변수로 노드 정보 node, 탐색 범위 s, e, 구간 일치 범위 L, R을 전달받는다.
2. 첫 번째 기저 조건으로 탐색 범위가 일치 범위를 벗어난 경우 0, 0을 리턴한다.
3. 두 번째 기저 조건으로 탐색 범위가 일치 범위 내에 있다면 현재 노드정보를 리턴한다.
4. 기저 조건에 해당하지 않는 경우 좌, 우 자식 노드로 재귀를 진행한다.
5. 재귀를 빠져나오며 각 값을 left, right변수에 저장한다.
6. left와 right의 구간합과 개수의 합을 더해 리턴해 준다.

 

문제풀이

1. n값을 입력 받고, long long타입의 벡터 results를 초기화 해준다.
2. n번의 반복문을 통해 심을 나무의 좌표를 변수 num에 입력 받아준다.
3. update 함수를 통해 num의 위치에 나무를 심어주고 세그먼트 트리를 최신화 해준다.
4. query 함수를 통해 0 ~ num - 1과 num + 1, N 범위를 탐색하고 결과를 각각 left, right로 받아준다.
5. 두 구간 모두 구간합 - num * 구간 개수 연산을 해준 후 결과 값을 더해 result변수에 저장해 준다.
6. 만약 첫번째 나무를 심은게 아니라면 results에 result % MOD한 값을 push해준다.
7. long long 타입 변수 ans를 1로 초기화 한 후 results에 저장된 값들을 곱해준 후 MOD로 나눈 나머지를 저장한다.
8. 최종적으로 ans에 저장된 값을 출력해 준다.

 

트러블 슈팅

1. 예제 4번의 답이 유독 출력값이 나오지 않았다, 확인해 보니 나무의 좌표가 중복되는 경우 처리를 잘못해 주었었다.
2. update시 이미 있는 노드에 더해주는 형식으로 변경해 주니 AC를 받았다.

 

참고 사항

• 나무를 심고, 해당 나무를 기준으로 좌, 우 구간합을 구하고, 구간의 개수만큼 현재 나무 위치 좌표를 빼주면 된다.
• 위 과정에서 음수가 나올 수 있으니 절대값을 해주는게 좋다(아마 오른쪽은 논리상 양수만 나올 것 같다.)

 

정답 코드

#include<iostream>
#include<vector>
#define ll long long
using namespace std;

const int N = 200000;
const int MOD = 1000000007;
int n;
pair<ll, int> tree[N * 4];

void update(int node, int s, int e, ll idx) {
	if (s == e) {
		tree[node].first += idx;
		tree[node].second++;
	}
	else {
		int mid = (s + e) / 2;
		if (idx <= mid) update(node * 2, s, mid, idx);
		else update(node * 2 + 1, mid + 1, e, idx);
		tree[node] = { tree[node * 2].first + tree[node * 2 + 1].first, tree[node * 2].second + tree[node * 2 + 1].second };
	}
}

pair<ll, int> query(int node, int s, int e, int L, int R) {
	if (R < s || e < L) return { 0, 0 };
	if (L <= s && e <= R) return tree[node];
	int mid = (s + e) / 2;
	pair<ll, int> left = query(node * 2, s, mid, L, R);
	pair<ll, int> right = query(node * 2 + 1, mid + 1, e, L, R);
	return { left.first + right.first, left.second + right.second };
}

int main() {
	ios::sync_with_stdio(0);
	cin.tie(0);

	cin >> n;
	vector<ll> results;
	for (int i = 0; i < n; ++i) {
		ll num; cin >> num;
		update(1, 0, N, num);
		pair<ll, int> left = query(1, 0, N, 0, num - 1);
		pair<ll, int> right = query(1, 0, N, num + 1, N);
		ll result = abs(left.first - num * left.second) + abs(right.first - num * right.second);
		if (i) results.push_back(result % MOD);
	}
	ll ans = 1;
	for (ll i : results) ans = (ans * i) % MOD;
	cout << ans;
}