[P5] 백준 14449번 Balanced Photo C++ 세그먼트 트리, 값/좌표 압축

리뷰

 

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

문제를 읽고 당연하게 update없이 최대값 트리를 build하고 query문을 통해 개수를 구한다고 생각했다.

하지만 바로 4%에서 시간초과를 맞게 되었고, 값의 범위가 10억이므로 값/좌표 압축이 필요하단 것을 알았다.

 

 

전역 변수

• N : 배열의 최대 크기를 정의할 상수 변수
• n : 배열의 크기를 저장할 변수
• lst : 초기 배열의 정보를 저장할 배열
• sorted : 정렬된 배열 상태를 저장할 배열
• comp : 정렬된 배열의 값을 key로, 인덱스를 value로 저장할 해시맵
• tree : 구간 합 세그먼트 트리 정보를 저장할 배열
• H : lst배열의 원소 순서를 기준으로 comp의 인덱스를 저장할 배열
• L : i번째 소의 왼쪽 소 중 자신보다 큰 소의 수를 저장할 배열
• R : i번째 소의 오른쪽 소 중 자신보다 큰 소의 수를 저장할 배열

 

함수

1. update

void update(int node, int s, int e, int idx) {
	if (s == e) tree[node]++;
	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] + tree[node * 2 + 1];
	}
}

 

세그먼트 트리 업데이트를 위한 함수

 

2. query

int query(int node, int s, int e, int L, int R) {
	if (R < s || e < L) return 0;
	if (L <= s && e <= R) return tree[node];
	int mid = (s + e) / 2;
	return query(node * 2, s, mid, L, R) + query(node * 2 + 1, mid + 1, e, L, R);
}

 

세그먼트 트리 구간합을 구하기 위한 함수

 

 

문제풀이

1. n값을 입력 받고, 0~n-1번째 요소를 lst배열에 입력 받고, sorted배열은 lst배열을 복사해 준다.
2. sort함수를 통해 sorted배열을 오름차순으로 정렬해 준다.
3. sorted배열의 0~n-1번째 요소를 순회하며 comp에 현재 요소의 값을 key로 인덱스를 value로 저장해 준다.
4. 0~n-1범위의 반복문을 추가로 개행하고, comp[lst[i]]를 H배열의 현재 인덱스 값으로 저장한다.
5. 0~n-1범위의 반복문을 또 개행한다, query함수를 통해 세그먼트 트리의 H[i]+1~n-1범위의 구간합을 L[i]에 저장해 준다.
6. update함수를 통해 세그먼트 트리의 H[i]번째 노드의 값을 1만큼 증가시켜 준다.
7. for문이 종료된 경우 memset함수를 통해 tree배열의 값을 모두 0으로 초기화 시켜준다.
8. n-1~0범위의 반복문을 개행하고, L[i]를 구할때와 동일하게 R[i]에도 query문의 결과값을 저장해 준다.
9. 변수 ans를 0으로 초기화 하고, 0~n-1번째 L, R배열의 요소를 순회해 준다.
10. L[i], R[i] 중 작은 값을 2배한 값이 L[i], R[i] 중 큰 값보다 작을 경우 ans를 증가시켜 준다.
11. ans에 저장된 값을 출력해 준다.

 

트러블 슈팅

1. 구간 최대값 트리 + 현재 요소의 왼쪽, 오른쪽에서 개수를 구하는 방식은 시간 초과가 유발되었다.
2. 값/좌표 압축을 진행하고 해당 인덱스를 토대로 구간 합쿼리 + 업데이트를 진행하여 AC를 받았다.

 

참고 사항

1. for문의 방향에 따라 현재 소보다 왼쪽, 오른쪽에서의 큰 소의 개수를 알 수 있게 된다.

 

정답 코드

#include<iostream>
#include<algorithm>
#include<unordered_map>
#include<cstring>
using namespace std;

const int N = 1e5;
int n;
int lst[N];
int sorted[N];
unordered_map<int, int> comp;
int tree[N * 4];
int H[N];
int L[N];
int R[N];

void update(int node, int s, int e, int idx) {
	if (s == e) tree[node]++;
	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] + tree[node * 2 + 1];
	}
}

int query(int node, int s, int e, int L, int R) {
	if (R < s || e < L) return 0;
	if (L <= s && e <= R) return tree[node];
	int mid = (s + e) / 2;
	return query(node * 2, s, mid, L, R) + query(node * 2 + 1, mid + 1, e, L, R);
}

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

	cin >> n;
	for (int i = 0; i < n; ++i) {
		cin >> lst[i];
		sorted[i] = lst[i];
	}

	sort(sorted, sorted + n);
	for (int i = 0; i < n; ++i) comp[sorted[i]] = i;

	for (int i = 0; i < n; ++i) H[i] = comp[lst[i]];

	for (int i = 0; i < n; ++i) {
		L[i] = query(1, 0, n - 1, H[i] + 1, n - 1);
		update(1, 0, n - 1, H[i]);
	}

	memset(tree, 0, sizeof(tree));
	for (int i = n - 1; i >= 0; --i) {
		R[i] = query(1, 0, n - 1, H[i] + 1, n - 1);
		update(1, 0, n - 1, H[i]);
	}

	int ans = 0;
	for (int i = 0; i < n; ++i) {
		if (min(L[i], R[i]) * 2 < max(L[i], R[i])) ans++;
	}
	cout << ans;
}