BOJ 7469 - K-th Number

On this page

Problem link

Why the original approach is too slow

Sorting the (value, index) pairs once takes O(N log N), but then scanning all N elements for each of the M range queries costs O(NM) in the worst case. Making the variables global rather than local changes storage duration, not the amount of work, so it does not change this asymptotic cost.

Persistent segment tree

Coordinate-compress the array values into sorted distinct ranks 0 ... U-1, where U is the number of distinct values. Build a persistent segment tree root for every array prefix: root[i] represents the first i elements. Each node stores the number of elements in its value interval. To add the next array value, copy the nodes on its root-to-leaf path and increment their counts, sharing every other node with the previous version.

For a query [L, R], the count of a value interval in the subarray is the count in root[R] minus the count in root[L-1]. Starting at the full rank interval, compare k with the difference between the two roots’ left-child counts. Descend left if the kth value is there; otherwise subtract that count from k and descend right. At a leaf, its rank identifies the kth smallest value. This also handles duplicates: each occurrence contributes one to the counts.

Building the compressed values and persistent roots takes O(N log U) time; each query takes O(log U). The nodes, input values, compressed coordinates, and prefix roots use O(N log U) space, and queries are processed one at a time without storing them.

C++17

#include <algorithm>
#include <iostream>
#include <vector>

using namespace std;

struct Node {
    int left = 0;
    int right = 0;
    int count = 0;
};


vector<Node> tree(1);

int update(int previous, int low, int high, int position) {
    const int current = static_cast<int>(tree.size());
    tree.push_back(tree[previous]);
    ++tree[current].count;

    if (low != high) {
        const int middle = low + (high - low) / 2;
        if (position <= middle) {
            tree[current].left = update(tree[previous].left, low, middle, position);
        } else {
            tree[current].right = update(tree[previous].right, middle + 1, high, position);
        }
    }
    return current;
}

int kth(int leftRoot, int rightRoot, int low, int high, int k) {
    if (low == high) return low;

    const int leftCount =
        tree[tree[rightRoot].left].count - tree[tree[leftRoot].left].count;
    const int middle = low + (high - low) / 2;

    if (k <= leftCount) {
        return kth(tree[leftRoot].left, tree[rightRoot].left, low, middle, k);
    }
    return kth(tree[leftRoot].right, tree[rightRoot].right, middle + 1, high,
               k - leftCount);
}

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

    int n, m;
    cin >> n >> m;

    vector<int> values(n);
    for (int& value : values) cin >> value;

    vector<int> coordinates = values;
    sort(coordinates.begin(), coordinates.end());
    coordinates.erase(unique(coordinates.begin(), coordinates.end()), coordinates.end());
    const int distinctCount = static_cast<int>(coordinates.size());

    vector<int> roots(n + 1, 0);
    for (int i = 0; i < n; ++i) {
        const int rank = static_cast<int>(
            lower_bound(coordinates.begin(), coordinates.end(), values[i]) - coordinates.begin());
        roots[i + 1] = update(roots[i], 0, distinctCount - 1, rank);
    }

    int left, right, k;
    while (m--) {
        cin >> left >> right >> k;
        const int rank = kth(roots[left - 1], roots[right],
                             0, distinctCount - 1, k);
        cout << coordinates[rank] << '\n';
    }
    return 0;
}