PS/C++

[백준 7469] K번째 수

siyamaki 2024. 6. 20. 11:19

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

 

머지 소트 트리 + find 시 binary search가 결합된 문제이다

 

처음 풀이에서는 find 시 binary search를 사용하지 않고 upper_bound로 했었는데 당연히 시간초과가 났었다

 

기본적인 풀이는 2024.06.20 - [PS/C++] - [백준 13544] 수열과 쿼리 3

 

[백준 13544] 수열과 쿼리 3

https://www.acmicpc.net/problem/13544 길이가 N인 수열 A1, A2, ..., AN이 주어진다. 이때, 다음 쿼리를 수행하는 프로그램을 작성하시오.i j k: Ai, Ai+1, ..., Aj로 이루어진 부분 수열 중에서 k보다 큰 원소의 개

download1324.tistory.com

의 풀이와 비슷하지만 K번째 수를 찾는 과정에서 binary search가 이용된다

 

정수의 값이 -1000000000 ~ 1000000000이니 이분 탐색을 + upper_bound를 이용해 구간 내 mid값을 조정 해 가면서 mid값보다 큰 수가 k보다 큰지, 작은지 확인해 가면서 start, end값을 조정해가면 된다

 

예를 들어 구간 내 값이 [1, 3, 4, 7, 9, 11]이고 5번째 수인 9를 찾으려고 한다.

시작값을 1, 끝 값을 15 라고 가정하면 

start end mid upper_bound
1 15 (1+15) = 16 / 2 = 8 4
upper bound가 k보다 작으니 start = mid + 1
9 15 (9 + 15) = 24 / 2 = 12 6
upper bound가 k보다 크니 end = mid - 1
9 11 (9 + 11) = 20 / 2 = 10 5
upper bound와 k가 같으니 end = mid - 1
9 9 9 5
upper bound와 k가 같으니 end = mid - 1
9 8 start <= end 조건에 맞지 않으므로 종료 후 start 반환

 

#include<bits/stdc++.h>
#include "vector"
#pragma GCC optimize ("O3")
#pragma GCC optimize ("Ofast")
#pragma GCC optimize ("unroll-loops")
#pragma GCC target("sse,sse2,sse3,ssse3,sse4,avx,avx2")

using namespace std;
int N, Q;
int arr[100001];
vector<int> segTree[262144];
void makeSegTree(int start, int end, int idx);
int find(int start, int end, int idx, int left, int right, int target);
int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr); cout.tie(nullptr);
    cin >> N >> Q;
    for(int i = 1; i <= N; i++) {
        cin >> arr[i];
    }
    makeSegTree(1, N, 1);
    while(Q-->0) {
        int l, r, t;
        cin >> l >> r >> t;
        int start = -1000000000;
        int end = 1000000000;
        while(start <= end) {
            int mid = (start + end) / 2;
            int res = find(1, N, 1, l, r, mid);
            if(res < t) {
                start = mid + 1;
            } else {
                end = mid - 1;
            }
        }
        cout << start << '\n';
    }
}
void makeSegTree(int start, int end, int idx) {
    if(start == end) {
        segTree[idx].push_back(arr[start]);
        return;
    }
    int mid = (end + start) / 2;
    makeSegTree(start, mid, idx << 1);
    makeSegTree(mid + 1, end, idx << 1 | 1);
    segTree[idx].insert(segTree[idx].end(), segTree[idx << 1].begin(), segTree[idx << 1].end());
    segTree[idx].insert(segTree[idx].end(), segTree[idx << 1 | 1].begin(), segTree[idx << 1 | 1].end());
    sort(segTree[idx].begin(), segTree[idx].end());
}
int find(int start, int end, int idx, int left, int right, int target) {    // target 이하인 수가 몇개인지 확인
    if(right < start || end < left) {
        return 0;
    }
    if(left <= start && end <= right) {
        return upper_bound(segTree[idx].begin(), segTree[idx].end(), target) - segTree[idx].begin();
    }
    int mid = (end + start) / 2;
    int l = find(start, mid, idx << 1, left, right, target);
    int r = find(mid + 1, end, idx << 1 | 1, left, right, target);
    return l + r;
}