kth_statistics

#include <iostream>
#include <fstream>
#include <vector>
#include <cstring>
#include <unordered_map>
#include <algorithm>
#include <deque>
#include <random>

using std::size_t;

const std::string TASK_NAME = "kth";

std::random_device rd;
std::mt19937 rng(rd());

int32_t find_kth(std::vector<int32_t> &vector, size_t k, size_t l = 0, size_t r = static_cast<size_t>(-1));

size_t split(std::vector<int32_t> &vector, size_t l, size_t r);

template<typename T>
constexpr T find_median(T x1, T x2, T x3);;

void solve(std::istream &in, std::ostream &out) {
    size_t n, k;
    in >> n >> k;
    --k;
    int32_t a, b, c;
    in >> a >> b >> c;

    std::vector<int32_t> xs(n);
    in >> xs[0];
    if (n != 1) {
        in >> xs[1];
    }
    for (std::size_t i = 2; i < n; ++i) {
        xs[i] = a * xs[i - 2] + b * xs[i - 1] + c;
    }

    out << find_kth(xs, k) << std::endl;
}

int main(int argc, char *argv[]) {
    bool is_local = argc >= 2 && strcmp(argv[1], "LOCAL") == 0;
    if (is_local) {
        std::ifstream fin("input.txt", std::ifstream::in);
        solve(fin, std::cout);
    } else {
        std::ifstream fin(TASK_NAME + ".in", std::ifstream::in);
        std::ofstream fout(TASK_NAME + ".out", std::ofstream::out);
        solve(fin, fout);
    }
    return 0;
}

int32_t find_kth(std::vector<int32_t> &vector, size_t k, size_t l, size_t r) {
    if (r == -1) {
        r = vector.size();
    }

    size_t m = split(vector, l, r);
    while (k != m) {
        if (k < m) {
            r = m;
        } else {
            l = m;
        }
        m = split(vector, l, r);
    }

    return *std::min_element(vector.begin() + m, vector.begin() + r);
}

size_t split(std::vector<int32_t> &vector, size_t l, size_t r) {
    if (r - l <= 1) {
        return l;
    }

    std::uniform_int_distribution<size_t> next_random(l, r - 1);

    size_t ind1 = next_random(rng);
    size_t ind2 = next_random(rng);
    size_t ind3 = next_random(rng);
    int32_t pivot = find_median(vector[ind1], vector[ind2], vector[ind3]);
    size_t lm = r - 1;
    size_t rm = l;
    while (rm <= lm) {
        while (vector[rm] < pivot) {
            ++rm;
        }
        while (vector[lm] > pivot) {
            --lm;
        }
        if (rm <= lm) {
            std::swap(vector[rm++], vector[lm]);
            lm = lm == 0 ? 0 : lm - 1;
        }
    }
    return lm + 1;
}

template<typename T>
constexpr T find_median(T x1, T x2, T x3) {
    if (std::min(x1, x2) != std::min(x2, x3)) {
        return std::max(std::min(x1, x2), std::min(x2, x3));
    } else {
        return std::min(x1, x3);
    }
}