[P2] 백준 7469번 K번째 수 C++ 세그먼트 트리, 이분 탐색, 머지 소트 트리

리뷰

 

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

6달만에 다시 들어와본 문제인데 쿼리문에서 merge없이 답을 도출할 방법을 고민해 보았다.

 

 

전역 변수

• N : 배열의 최대 크기를 저장할 상수 변수
• lst : 요소의 정보를 저장할 배열
• tree : 세그먼트 트리 정보를 저장할 벡터 배열
• n : 배열의 크기를 저장할 변수
• m : 쿼리의 개수를 저장할 변수

 

함수

1. build

void build(int node, int s, int e)

 

세그먼트 트리 정보를 초기화 하기 위한 함수

1. 매개 변수로 노드 정보 node, 탐색 범위 s, e를 전달 받는다.
2. 리프 노드에 도달했을 경우 각 요소의 값으로 초기화 해준다.
3. 좌, 우 자식 노드로 재귀를 호출하며 두 노드를 merge하여 현재 노드에 정렬된 상태로 저장한다.

 

2. bs

int bs(int s, int e, int k)

 

이진 탐색을 통해 정렬된 구간에서 k번째 수를 구하는 함수

1. 매개 변수로 탐색 범위 s, e와 뽑아낼 수가 몇번째인지를 알 k를 전달 받는다.
2. 정수의 범위는 절댓값 10^9를 넘지 않으므로 L, R을 -1^9, 1^9로 초기화 해준다.
3. 이진 탐색을 통해 현재 mid값을 query함수에 전달하여 mid보다 작거나 같은 개수를 변수 cnt에 저장해 준다.
4. cnt가 k이상일 경우 R범위를 줄이고, 아닐 경우 L범위를 줄여준다.
5. while 루프가 종료된 후엔 L에 저장된 값을 출력해 준다.

 

3. query

int query(int node, int s, int e, int L, int R, int val)

 

쿼리 구간에서 val이하인 값의 개수를 구하는 함수

1. 매개 변수로 노드 정보 node, 탐색 범위 s, e 쿼리 구간 L, R 기준 값 val을 전달 받는다.
2. 쿼리 구간이 탐색 범위를 벗어난 경우 0을 리턴해 준다.
3. 탐색 범위가 쿼리 구간과 일치할 경우 upper_bound를 통해 val이하인 원소 개수를 리턴해 준다.
4. 위 기저 조건에 해당하지 않는 경우 좌, 우 자식 노드로 재귀를 진행하고, 각 노드의 리턴 값의 합을 리턴한다.

 

문제풀이

1. n, m값을 입력 받고, lst배열에 n개의 원소를 입력 받아준다.
2. build 함수를 통해 세그먼트 트리 정보를 초기화 해준다.
3. m번의 while 루프를 수행하고, 매 루프마다 인자 i, j, k를 입력 받아 bs함수에 매개변수로 전달해 준다.
4. 각 쿼리마다 결과값을 출력해 주고 줄 바꿈을 기준으로 구분해 준다.

 

트러블 슈팅

1. query 함수 내에서 merge를 돌린다면 시간 초과가 발생한다.
2. query 함수 내에서 k를 인자로 전달하고 범위에 따라 k의 개수를 줄여주는 풀이는 최적해가 나올 수 없다.
3. 이진 탐색을 통해 val값을 조절해 가며 val이하인 개수가 k개인 경우를 찾으면 최적해를 찾을 수 있다.

 

참고 사항

없음

 

 

정답 코드

#include<iostream>
#include<vector>
#include<algorithm>
using namespace std;
const int N = 1e5 + 1;
int lst[N];
vector<int> tree[N * 4];
int n, m;

void build(int node, int s, int e) {
	if (s == e) tree[node] = { lst[s] };
	else {
		int mid = (s + e) / 2;
		build(node * 2, s, mid);
		build(node * 2 + 1, mid + 1, e);
		vector<int>& L = tree[node * 2];
		vector<int>& R = tree[node * 2 + 1];
		tree[node].resize(L.size() + R.size());
		merge(L.begin(), L.end(), R.begin(), R.end(), tree[node].begin());
	}
}

int query(int node, int s, int e, int L, int R, int val) {
	if (R < s || e < L) return 0;
	if (L <= s && e <= R) return upper_bound(tree[node].begin(), tree[node].end(), val) - tree[node].begin();
	int mid = (s + e) / 2;
	int left = query(node * 2, s, mid, L, R, val);
	int right = query(node * 2 + 1, mid + 1, e, L, R, val);
	return left + right;
}

int bs(int s, int e, int k) {
	int L = -1e9, R = 1e9;
	while (L <= R) {
		int mid = (L + R) / 2;
		int cnt = query(1, 1, n, s, e, mid);
		if (cnt >= k) R = mid - 1;
		else L = mid + 1;
	}
	return L;
}

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

	cin >> n >> m;
	for (int i = 1; i <= n; ++i) cin >> lst[i];

	build(1, 1, n);
	while (m--) {
		int i, j, k; cin >> i >> j >> k;
		cout << bs(i, j, k) << "\n";
	}
}