Implementing 3 way partitioning for quick sort

Question

I am trying to implement 3-Way partitioning for quick sort. I tested with a lot of custom test cases but if works fine, but it fails for some unknown cases. I am unable to figure out where am I going.

Here's the code. Partitioning:

int* partition3(vector<long long int> &a, int l, int r) {
  int* m = new int[2];

  long long int x = a[l];

  int j = l;
  int k = l;

  for (int i = l + 1; i <= r; i++) {

    if (a[i] < x) {
      j++;
      k++;
      swap(a[i], a[j]);
    }
    else if(a[i]==x)
    { 
      k++;
      swap(a[i],a[k]);
    }
  }
  swap(a[l], a[j]);    
  m[0]=j;
  m[1]=k;
  return m;
}

Sort Function:

void randomized_quick_sort(vector<long long int> &a, int l, int r) {
  if (l >= r) {
    return;
  }

  int k = l + rand() % (r - l + 1);

  swap(a[l], a[k]);

  int* m = new int[2];

  m = partition3(a, l, r);

  randomized_quick_sort(a, l, m[0]-1);
  randomized_quick_sort(a, m[1], r);
}

I will be grateful if you help me out.

Answer 1

The easiest way to implement three-way partition that is provably correct is to use the method of loop invariants. For simplicity and genericity let's work with iterators. Consider the following invariants:

In the range
  [first, i)  all elements are less than pivot
  [i, j)      all elements are equal to pivot
  [j, k)      unpartitioned range
  [k, last)   all elements are greater than pivot

Initially, i = first, j = first, and k = last, so that the whole range [first, last) is unpartitioned. At each iteration we'll contract this range by one element. Finally, j = k, so that the whole range is three-way partitioned.

The following code implements this idea:

template<class It>
std::pair<It, It> three_way_partition(It first, It last) {
    assert(first != last);    
    const auto& pivot = *--last;

    auto i = first, j = first, k = last;
    while (j != k)
        if (*j < pivot)
            std::iter_swap(i++, j++);
        else if (pivot < *j)
            std::iter_swap(j, --k);
        else
            ++j;

    std::iter_swap(j++, last);
    return {i, j};
}

Here I used the last element as the pivot. This choice simplifies the code, but is not essential.

Quick sort algorithm that uses this function is:

template<class It, class Gen>
void randomized_quick_sort(It first, It last, Gen&& gen) {
    if (last - first <= 1)
        return;

    std::uniform_int_distribution<typename It::difference_type> 
        dist(0, last - first - 1);
    std::iter_swap(first + dist(gen), last - 1);

    const auto p = three_way_partition(first, last);
    randomized_quick_sort(first, p.first, gen);
    randomized_quick_sort(p.second, last, gen);
}

Sample use case:

std::mt19937 gen; // you might want to initialize it with std::random_device
std::vector<int> vec;
// initialize vec

randomized_quick_sort(vec.begin(), vec.end(), gen);

Demo