[Numpy-svn] r4802 - in trunk/numpy/core: src tests

[numpy-svn at scipy.org]

Author: charris
Date: 2008-02-13 17:10:44 -0600 (Wed, 13 Feb 2008)
New Revision: 4802

Modified:
   trunk/numpy/core/src/_sortmodule.c.src
   trunk/numpy/core/tests/test_multiarray.py
Log:
Add type specific heapsort for strings and unicode.


Modified: trunk/numpy/core/src/_sortmodule.c.src
===================================================================
--- trunk/numpy/core/src/_sortmodule.c.src	2008-02-13 21:41:47 UTC (rev 4801)
+++ trunk/numpy/core/src/_sortmodule.c.src	2008-02-13 23:10:44 UTC (rev 4802)
@@ -520,10 +520,57 @@
 
 
 static int
+ at TYPE@_heapsort(@type@ *start, intp n, PyArrayObject *arr)
+{
+    size_t len = arr->descr->elsize/sizeof(@type@);
+    @type@ *tmp = malloc(arr->descr->elsize);
+    @type@ *a = start - len;
+    intp i,j,l;
+
+    for (l = n>>1; l > 0; --l) {
+        @copy@(tmp, a + l*len, len);
+        for (i = l, j = l<<1; j <= n;) {
+            if (j < n && @lessthan@(a + j*len, a + (j+1)*len, len))
+                j += 1;
+            if (@lessthan@(tmp, a + j*len, len)) {
+                @copy@(a + i*len, a + j*len, len);
+                i = j;
+                j += j;
+            }
+            else
+                break;
+        }
+        @copy@(a + i*len, tmp, len);
+    }
+
+    for (; n > 1;) {
+        @copy@(tmp, a + n*len, len);
+        @copy@(a + n*len, a + len, len);
+        n -= 1;
+        for (i = 1, j = 2; j <= n;) {
+            if (j < n && @lessthan@(a + j*len, a + (j+1)*len, len))
+                j++;
+            if (@lessthan@(tmp, a + j*len, len)) {
+                @copy@(a + i*len, a + j*len, len);
+                i = j;
+                j += j;
+            }
+            else
+                break;
+        }
+        @copy@(a + i*len, tmp, len);
+    }
+
+    free(tmp);
+    return 0;
+}
+
+
+static int
 @TYPE at _aheapsort(@type@ *v, intp *tosort, intp n, PyArrayObject *arr)
 {
+    size_t len = arr->descr->elsize/sizeof(@type@);
     intp *a, i,j,l, tmp;
-    size_t len = arr->descr->elsize/sizeof(@type@);
 
     /* The array needs to be offset by one for heapsort indexing */
     a = tosort - 1;
@@ -569,13 +616,13 @@
 static int
 @TYPE at _aquicksort(@type@ *v, intp* tosort, intp num, PyArrayObject *arr)
 {
+    size_t len = arr->descr->elsize/sizeof(@type@);
     @type@ *vp;
     intp *pl = tosort;
     intp *pr = tosort + num - 1;
     intp *stack[PYA_QS_STACK];
     intp **sptr=stack;
     intp *pm, *pi, *pj, *pk, vi, SWAP_temp;
-    size_t len = arr->descr->elsize/sizeof(@type@);
 
     for(;;) {
         while ((pr - pl) > SMALL_QUICKSORT) {
@@ -725,6 +772,8 @@
         (PyArray_ArgSortFunc *)STRING_aheapsort;
     descr->f->sort[PyArray_QUICKSORT] = \
         (PyArray_SortFunc *)STRING_quicksort;
+    descr->f->sort[PyArray_HEAPSORT] = \
+        (PyArray_SortFunc *)STRING_heapsort;
 
     descr = PyArray_DescrFromType(PyArray_UNICODE);
     descr->f->argsort[PyArray_MERGESORT] = \
@@ -735,6 +784,8 @@
         (PyArray_ArgSortFunc *)UNICODE_aheapsort;
     descr->f->sort[PyArray_QUICKSORT] = \
         (PyArray_SortFunc *)UNICODE_quicksort;
+    descr->f->sort[PyArray_HEAPSORT] = \
+        (PyArray_SortFunc *)UNICODE_heapsort;
 }
 
 static struct PyMethodDef methods[] = {

Modified: trunk/numpy/core/tests/test_multiarray.py
===================================================================
--- trunk/numpy/core/tests/test_multiarray.py	2008-02-13 21:41:47 UTC (rev 4801)
+++ trunk/numpy/core/tests/test_multiarray.py	2008-02-13 23:10:44 UTC (rev 4802)
@@ -290,11 +290,11 @@
             c.sort(kind=kind)
             assert_equal(c, ai, msg)
 
-        # test string sorts. Only quicksort is available at this time.
+        # test string sorts. Only quick and heap sort are available.
         s = 'aaaaaaaa'
         a = np.array([s + chr(i) for i in range(100)])
         b = a[::-1].copy()
-        for kind in ['q'] :
+        for kind in ['q', 'h'] :
             msg = "string sort, kind=%s" % kind
             c = a.copy();
             c.sort(kind=kind)
@@ -303,11 +303,11 @@
             c.sort(kind=kind)
             assert_equal(c, a, msg)
 
-        # test unicode sort. Only quicksort is available at this time.
+        # test unicode sort. Only quick and heap sort are available.
         s = 'aaaaaaaa'
         a = np.array([s + chr(i) for i in range(100)], dtype=np.unicode)
         b = a[::-1].copy()
-        for kind in ['q'] :
+        for kind in ['q', 'h'] :
             msg = "unicode sort, kind=%s" % kind
             c = a.copy();
             c.sort(kind=kind)