k-smallest selection using max-heap

k-smallest.c

/* Find K-smallest value in array use a simplified max-heap */

/*
  0. Build a max-heap(MH) for the first K elements
  1. For each element after the Kth element, compare it with the root of MH
     0. If the element is smaller than the root, make it root and heapify^1 MH
     1. Else ignore it.
     2. Finally, MH has K-smallest elements.

  Complexity: O(K + (n-K)log(K) + Klog(K)), the last Klog(K) for sorted^2 ouput.

  [1]: heapify: Here just find the max value and swap to the root.
  [2]: Here I used a 2-insertion sort, 'cause the data length was considered small.
*/

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

struct maxheap {
    int n;
    double *data;
};

#define maxheap_root(mh) ((mh)->data[0])

static inline int
_maxheap_findmax(const double *arr, int n) {
    int idx = 0;
    double t = arr[0];
    for (int i = 0; i < n; ++i) {
	if (t < arr[i]) {
	    t = arr[i];
	    idx = i;
	}
    }
    return idx;
}

static inline void
_maxheap_swap(double *arr, int i, int j) {
    double t = arr[i];
    arr[i] = arr[j];
    arr[j] = t;
}

struct maxheap *
maxheap_new(const double *arr, int n) {
    int idx = 0;
    struct maxheap *mh= (struct maxheap*) malloc ( sizeof(*mh) + sizeof(double) * n );
    mh->n = n;
    mh->data = (double*)(mh + 1);

    memcpy(mh->data, arr, sizeof(double) * n);
    idx = _maxheap_findmax(mh->data, n);
    _maxheap_swap(mh->data, 0, idx);
    return mh;
}

void
maxheap_free(struct maxheap **mh) {
    if (mh) {
	free ( *mh );
	*mh = 0;
    }
}

void
maxheap_insert(struct maxheap *mh, double small) {
    int idx;
    mh->data[0] = small;
    idx = _maxheap_findmax(mh->data, mh->n);
    _maxheap_swap(mh->data, 0, idx);
}

/* 2-insertion sort from "Sorting Methods for small arrays" Renato F. Werneck, el.*/
static inline void
_maxheap_sort(double *v, int left, int right) {
    double tmin, tmax;
    int j, i = (left + (right - left + 1)) % 2;
    
    for (; i < right; ++i) {
	tmin = v[i];
	tmax = v[i+1];

	if ( v[i] > v[i+1] ) {
	    tmax = v[i];
	    tmin = v[i+1];
	}

	j = i - 1;
	while (v[j] > tmax) {
	    v[j + 2] = v[j];
	    j--;
	}
	v[j+2] = tmax;

	while (v[j] > tmin) {
	    v[j+1] = v[j];
	    j--;
	}
	v[j+1] = tmin;
    }    
}

void
maxheap_sort(struct maxheap *mh) {
    _maxheap_sort(mh->data, 0, mh->n - 1);
}

/* select K-smallest elements in an array */
void
select_ksmallest(const double *arr, int n, int k, double *out) {
    struct maxheap *mh = maxheap_new(arr, k);
    for (int i = k; i < n; ++i) {
	if ( arr[i] < maxheap_root(mh) ) {
	    maxheap_insert(mh, arr[i]);
	}
    }

    maxheap_sort(mh);
    memcpy(out, mh->data, sizeof(double) * k);    
    maxheap_free( &mh );
}

/* unit test */
static void 
printv(double *v, int n) {
    for (int i = 0; i < n; ++i) {
	printf("%lf ", v[i]);
    }
}

int main(int argc, char *argv[]) {
    double arr[10] = {1,3,4,5,6,7,3,1,4,5};
    double out[10];
    for (int i = 1; i <= 10; ++i) {
	select_ksmallest(arr, 10, i, out);
	printv(out, i);
	printf("\n");
    }
    getchar();
    return 0;
}