// Selección: el k-ésimo menor de un arreglo, con k contado desde 0.
// `comparisons` cuenta cada vez que un elemento se compara con un pivote o, dentro de un
// grupo, con otro elemento. Es un conteo exacto: no depende de la máquina.
#ifndef SELECT_H
#define SELECT_H

#include <algorithm>
#include <cassert>
#include <cstddef>
#include <random>
#include <utility>
#include <vector>

struct Parts {
    std::vector<int> less;
    std::vector<int> equal;
    std::vector<int> greater;
};

// Parte en tres: menores, iguales y mayores que el pivote. Cuesta una comparación por elemento.
inline Parts partition3(const std::vector<int>& a, int pivot, long long& comparisons) {
    Parts parts;
    for (int x : a) {
        comparisons++;
        if (x < pivot) {
            parts.less.push_back(x);
        } else if (x > pivot) {
            parts.greater.push_back(x);
        } else {
            parts.equal.push_back(x);
        }
    }
    return parts;
}

// Quickselect con una regla para elegir el pivote.
// Invariante: la respuesta buscada es el k-ésimo menor de `a`.
template <typename PivotRule>
int quickselect(std::vector<int> a, std::size_t k, PivotRule choose_pivot, long long& comparisons) {
    assert(k < a.size());
    while (a.size() > 1) {
        Parts parts = partition3(a, choose_pivot(a), comparisons);
        if (k < parts.less.size()) {
            a = std::move(parts.less);
        } else if (k < parts.less.size() + parts.equal.size()) {
            return parts.equal[0];
        } else {
            k -= parts.less.size() + parts.equal.size();
            a = std::move(parts.greater);
        }
    }
    return a[0];
}

inline int quickselect_first(std::vector<int> a, std::size_t k, long long& comparisons) {
    return quickselect(std::move(a), k, [](const std::vector<int>& v) { return v[0]; }, comparisons);
}

inline int quickselect_random(std::vector<int> a, std::size_t k, std::mt19937& rng, long long& comparisons) {
    auto rule = [&rng](const std::vector<int>& v) { return v[rng() % v.size()]; };
    return quickselect(std::move(a), k, rule, comparisons);
}

// Ordena un grupo pequeño por inserción. En un grupo de 5 son a lo sumo 10 comparaciones.
inline void sort_group(std::vector<int>& group, long long& comparisons) {
    for (std::size_t i = 1; i < group.size(); i++) {
        int value = group[i];
        std::size_t j = i;
        while (j > 0) {
            comparisons++;
            if (group[j - 1] <= value) {
                break;
            }
            group[j] = group[j - 1];
            j--;
        }
        group[j] = value;
    }
}

// Reparte en grupos de 5 por posición y devuelve la mediana inferior de cada uno. El último grupo puede ser más corto.
inline std::vector<int> group_medians(const std::vector<int>& a, long long& comparisons) {
    std::vector<int> medians;
    for (std::size_t start = 0; start < a.size(); start += 5) {
        std::size_t stop = std::min(start + 5, a.size());
        std::vector<int> group(a.begin() + start, a.begin() + stop);
        sort_group(group, comparisons);
        medians.push_back(group[(group.size() - 1) / 2]);
    }
    return medians;
}

int select_linear(std::vector<int> a, std::size_t k, long long& comparisons);

// El pivote de la mediana de medianas: la mediana inferior de las medianas de los grupos.
inline int median_of_medians(const std::vector<int>& a, long long& comparisons) {
    std::vector<int> medians = group_medians(a, comparisons);
    std::size_t middle = (medians.size() - 1) / 2;
    return select_linear(std::move(medians), middle, comparisons);
}

// Selección en tiempo lineal en el peor caso.
inline int select_linear(std::vector<int> a, std::size_t k, long long& comparisons) {
    assert(k < a.size());
    if (a.size() <= 5) {
        sort_group(a, comparisons);
        return a[k];
    }
    int pivot = median_of_medians(a, comparisons);
    Parts parts = partition3(a, pivot, comparisons);
    if (k < parts.less.size()) {
        return select_linear(std::move(parts.less), k, comparisons);
    }
    if (k < parts.less.size() + parts.equal.size()) {
        return pivot;
    }
    return select_linear(std::move(parts.greater), k - parts.less.size() - parts.equal.size(), comparisons);
}

// Una entrada que lleva la mediana de medianas casi hasta su cota: n múltiplo de 5, todos distintos.
// La mitad de los grupos (redondeada hacia arriba) tiene sus tres menores entre los valores
// más pequeños de todos, y el resto de los valores queda por encima del pivote.
inline std::vector<int> tight_input(int n) {
    assert(n % 5 == 0);
    int groups = n / 5;
    int low_groups = (groups + 1) / 2;
    int next_low = 0;
    int next_high = 3 * low_groups;
    std::vector<int> a;
    for (int g = 0; g < groups; g++) {
        int lows = g < low_groups ? 3 : 0;
        for (int i = 0; i < lows; i++) {
            a.push_back(next_low++);
        }
        for (int i = lows; i < 5; i++) {
            a.push_back(next_high++);
        }
    }
    return a;
}

#endif