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

numpy-svn@scip... numpy-svn@scip...
Wed Feb 13 15:41:57 CST 2008


Author: charris
Date: 2008-02-13 15:41:47 -0600 (Wed, 13 Feb 2008)
New Revision: 4801

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


Modified: trunk/numpy/core/src/_sortmodule.c.src
===================================================================
--- trunk/numpy/core/src/_sortmodule.c.src	2008-02-13 20:37:28 UTC (rev 4800)
+++ trunk/numpy/core/src/_sortmodule.c.src	2008-02-13 21:41:47 UTC (rev 4801)
@@ -520,6 +520,53 @@
 
 
 static int
+@TYPE@_aheapsort(@type@ *v, intp *tosort, intp n, PyArrayObject *arr)
+{
+    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;
+
+    for (l = n>>1; l > 0; --l) {
+        tmp = a[l];
+        for (i = l, j = l<<1; j <= n;) {
+            if (j < n && @lessthan@(v + a[j]*len, v + a[j+1]*len, len))
+                j += 1;
+            if (@lessthan@(v + tmp*len, v + a[j]*len, len)) {
+                a[i] = a[j];
+                i = j;
+                j += j;
+            }
+            else
+                break;
+        }
+        a[i] = tmp;
+    }
+
+    for (; n > 1;) {
+        tmp = a[n];
+        a[n] = a[1];
+        n -= 1;
+        for (i = 1, j = 2; j <= n;) {
+            if (j < n && @lessthan@(v + a[j]*len, v + a[j+1]*len, len))
+                j++;
+            if (@lessthan@(v + tmp*len, v + a[j]*len, len)) {
+                a[i] = a[j];
+                i = j;
+                j += j;
+            }
+            else
+                break;
+        }
+        a[i] = tmp;
+    }
+
+    return 0;
+}
+
+
+static int
 @TYPE@_aquicksort(@type@ *v, intp* tosort, intp num, PyArrayObject *arr)
 {
     @type@ *vp;
@@ -674,6 +721,8 @@
         (PyArray_ArgSortFunc *)STRING_amergesort;
     descr->f->argsort[PyArray_QUICKSORT] = \
         (PyArray_ArgSortFunc *)STRING_aquicksort;
+    descr->f->argsort[PyArray_HEAPSORT] = \
+        (PyArray_ArgSortFunc *)STRING_aheapsort;
     descr->f->sort[PyArray_QUICKSORT] = \
         (PyArray_SortFunc *)STRING_quicksort;
 
@@ -682,6 +731,8 @@
         (PyArray_ArgSortFunc *)UNICODE_amergesort;
     descr->f->argsort[PyArray_QUICKSORT] = \
         (PyArray_ArgSortFunc *)UNICODE_aquicksort;
+    descr->f->argsort[PyArray_HEAPSORT] = \
+        (PyArray_ArgSortFunc *)UNICODE_aheapsort;
     descr->f->sort[PyArray_QUICKSORT] = \
         (PyArray_SortFunc *)UNICODE_quicksort;
 }

Modified: trunk/numpy/core/tests/test_multiarray.py
===================================================================
--- trunk/numpy/core/tests/test_multiarray.py	2008-02-13 20:37:28 UTC (rev 4800)
+++ trunk/numpy/core/tests/test_multiarray.py	2008-02-13 21:41:47 UTC (rev 4801)
@@ -365,26 +365,24 @@
             assert_equal(ai.copy().argsort(kind=kind), a, msg)
             assert_equal(bi.copy().argsort(kind=kind), b, msg)
 
-        # test string argsorts. Only quick and merge sort are
-        # available at this time.
+        # test string argsorts.
         s = 'aaaaaaaa'
         a = np.array([s + chr(i) for i in range(100)])
         b = a[::-1].copy()
         r = arange(100)
         rr = r[::-1].copy()
-        for kind in ['q', 'm'] :
+        for kind in ['q', 'm', 'h'] :
             msg = "string argsort, kind=%s" % kind
             assert_equal(a.copy().argsort(kind=kind), r, msg)
             assert_equal(b.copy().argsort(kind=kind), rr, msg)
 
-        # test unicode argsorts. Only quick and merge sort are
-        # available at this time.
+        # test unicode argsorts.
         s = 'aaaaaaaa'
         a = np.array([s + chr(i) for i in range(100)], dtype=np.unicode)
         b = a[::-1].copy()
         r = arange(100)
         rr = r[::-1].copy()
-        for kind in ['q', 'm'] :
+        for kind in ['q', 'm', 'h'] :
             msg = "unicode argsort, kind=%s" % kind
             assert_equal(a.copy().argsort(kind=kind), r, msg)
             assert_equal(b.copy().argsort(kind=kind), rr, msg)



More information about the Numpy-svn mailing list