find-median.cpp

#include <iostream>

using namespace std;

void swap(int A[], int i, int j)
{
    int temp = A[i];
    A[i] = A[j];
    A[j] = temp;
}

/**
 * Partition the array in [low..high].
 * @return int index of the pivot
 */
int partition(int A[], int low, int high)
{
    int pivot = A[high];
    int i = low, j = high;
    while(true) {
        while(i < j && A[i] < pivot) ++i;
        while(i < j && A[j] >= pivot) --j;
        
        if(i == j) {
            // put pivot at boudnary
            swap(A, i, high);
            return i;
        }
        swap(A, i, j);
    }
}

/**
 * Find median of n unsorted numbers.
 * NOTE: assuming n > 0 && n%2 == 1
 * complexity: O(n)
 */
int findMedian(int A[], int n)
{
    int low = 0, high = n-1;
    int pivot_idx, center = n / 2;
    
    while(true) {
        pivot_idx = partition(A, low, high);
        if(pivot_idx == center)
            return A[pivot_idx];
        else if(pivot_idx < center)
            low = pivot_idx+1;
        else
            high = pivot_idx-1;
   }
}

/**
 * Test cases:
 * 1. length of array (assume odd in our case)
 *    a. empty array: []
 *    b. odd length: [3, 1, 2] -> 2
 *    c. even length: [2, 3, 1 4] -> 2 or 3?
 * 2. the sortedness of array
 *    a. already sorted: [1, 2, 3] -> 2
 *    b. unsorted: [2, 3, 4, 5, 1] -> 3
 * 3. content:
 *    a. same number: [2, 2, 2, 2] -> 2
 *    b. mix positive with negative: [-3, 2, 0, -1, 1] -> 0
 */

int main() {
    int A[] = {3, 1, 2};
    cout << (2 == findMedian(A, 3)) << endl;
    int B[] = {2, 3, 4, 5, 1};
    cout << (3 == findMedian(B, 5)) << endl;
    int C[] = {1, 2, 3};
    cout << (2 == findMedian(C, 3)) << endl;
    int D[] = {2, 2, 2, 2, 2};
    cout << (2 == findMedian(D, 5)) << endl;

    return 0;
}