Алгоритм нахождения медианы

Реализую алгоритм нахождения медианы, при этом массив предварительно разбиваю относительно заданного опорного элемента pivot на 3 части - подмассив элементов, меньших pivot, подмассив элементов, равных pivot, и подмассив элементов, бОльших pivot. Изначально на вход подаются размер массива, его элементы и индекс опорного элемента (в дальнейшем он генерируется внутри функции findMedian). Данный алгоритм не работает на тесте с такими массивом:

4 3 5 7 -2 8 19 20 9

Если изначально задавать разбиение относительно элемента с индексом 0 (то есть 4), то медиана равна 9. Если же взять индекс 7 (20), то медиана равна 8, а если индекс равен 4 (-2), то медиана равна 7. Помогите разобраться, в чем ошибка :( (возможно я совсем неправильно понимаю суть алгоритма поиска k-ого по величине элемента)

#include <iostream>
#include <vector>
#include <algorithm>
#include <utility>
#include <random>

std::pair<std::size_t, std::size_t> split(std::vector<int>&, std::size_t, std::size_t, std::size_t);
int randomPivot(std::size_t, std::size_t);
int findMedian(std::vector<int>&, std::size_t, std::size_t, int);

int main()
{
  std::vector<int> vector {};
  std::size_t size = 0;
  int element = 0;
  int pivot = 0;

  std::cout << "Enter the size of the array: ";
  std::cin >> size;
  std::cout << "\nEnter array elements: ";

  while (size != 0)
  {
    std::cin >> element;
    vector.push_back(element);
    size--;
  }
  size = vector.size();
  
  std::cout << "\nEnter the index of pivot: ";
  std::cin >> pivot;

  int k = (size % 2 == 0 ? (size / 2) - 1 : size / 2);
  int median = 0;

  median = findMedian(vector, size, k, pivot);
  std::cout << "\nMedian: " << median;

  return 0;
}

std::pair<std::size_t, std::size_t> split(std::vector<int>& array, std::size_t begin, std::size_t end, std::size_t index)
{
  int pivot = array[index];
  std::size_t temp = 0;

  for (std::size_t i = begin; i < end; ++i)
  {
    if (array[i] < pivot)
    {
      std::swap(array[i], array[temp]);
      temp++;
    }
  }

  std::size_t minIndex = temp;
  std::size_t maxIndex = 0;

  for (std::size_t i = temp; i < end; ++i)
  {
    if (array[i] == pivot)
    {
      std::swap(array[i], array[temp]);
      maxIndex = temp;
      temp++;
    }
  }

  std::pair<std::size_t, std::size_t> indices(minIndex, maxIndex);

  return indices;
}

int findMedian(std::vector<int>& vector, std::size_t size, std::size_t k, int pivot)
{
  if (size == 1)
  {
    return vector[0];
  }

  std::size_t left = 0;
  std::size_t right = size;

  std::pair<std::size_t, std::size_t> minAndMax = split(vector, left, right, pivot);

  int median = 0;

  if ((k >= minAndMax.first) && (k <= minAndMax.second))
  {
    median = vector[minAndMax.first];
  }
  else if (k < minAndMax.first)
  {
    std::vector<int> leftArray;

    for (std::size_t i = 0; i < minAndMax.first; ++i)
    {
      leftArray.push_back(vector[i]);
    }
    pivot = randomPivot(0, leftArray.size() - 1);
    median = findMedian(leftArray, minAndMax.first, k, pivot);
  }
  else if (k > minAndMax.second)
  {
    std::vector<int> rightArray;

    for (std::size_t i = minAndMax.second + 1; i < size; ++i)
    {
      rightArray.push_back(vector[i]);
    }
    pivot = randomPivot(0, rightArray.size() - 1);
    median = findMedian(rightArray, rightArray.size(), k - minAndMax.second - 1, pivot);
  }
  
  return median;
}

int randomPivot(std::size_t left, std::size_t right)
{
  std::random_device random_device;
  std::mt19937 generator(random_device());

  std::uniform_int_distribution<> distribution(left, right);

  int x = distribution(generator);

  return x;
}

Ответы (1 шт):

Автор решения: Mikhailo

Если просто поиск k-го элемента, то вот:

int n_th(vector<int>& v, int k, int start, int end) {
    if (start == end) return v[k];   // Только один элемент

    if (k < start || k > end) throw runtime_error("Wrong k in n_th");
    // Разделение 
    int i = start, j = end, pivot = v[rand() % (j - i + 1) + i];

    for (; i < j;) {
        while (i < j && v[i] < pivot) ++i;

        while (i < j && v[j] >= pivot) --j;

        if (i < j) {
            int t = v[i];
            v[i] = v[j];
            v[j] = t;
            }
        }
    // Разделили. Смотрим, в какой части k и работаем только с ней
    if (k >= j) return n_th(v, k, j, end);
    else return n_th(v, k, start, j - 1);
    }
→ Ссылка