Merge pull request #57 from forFudan/zyharray

[lib] Implement the `argsort` functionality

Author: MadAlex1997

numojo/core/sort.mojo

@@ -155,3 +155,99 @@ fn binary_sort[
                 result[i - 1] = result[i]
                 result[i] = temp
     return result
+
+
+# ===------------------------------------------------------------------------===#
+# Argsort using quick sort algorithm
+# ===------------------------------------------------------------------------===#
+
+
+fn _argsort_partition(
+    inout ndarray: NDArray,
+    inout idx_array: NDArray,
+    left: Int,
+    right: Int,
+    pivot_index: Int,
+) raises -> Int:
+    var pivot_value = ndarray[pivot_index]
+    ndarray[pivot_index], ndarray[right] = ndarray[right], ndarray[pivot_index]
+    idx_array[pivot_index], idx_array[right] = (
+        idx_array[right],
+        idx_array[pivot_index],
+    )
+    var store_index = left
+
+    for i in range(left, right):
+        if ndarray[i] < pivot_value:
+            ndarray[store_index], ndarray[i] = ndarray[i], ndarray[store_index]
+            idx_array[store_index], idx_array[i] = (
+                idx_array[i],
+                idx_array[store_index],
+            )
+            store_index = store_index + 1
+
+    ndarray[right], ndarray[store_index] = ndarray[store_index], ndarray[right]
+    idx_array[right], idx_array[store_index] = (
+        idx_array[store_index],
+        idx_array[right],
+    )
+
+    return store_index
+
+
+fn argsort_inplace[
+    dtype: DType
+](
+    inout ndarray: NDArray[dtype],
+    inout idx_array: NDArray[DType.index],
+    left: Int,
+    right: Int,
+) raises:
+    """
+    Conduct Argsort (in-place) based on the NDArray using quick sort.
+
+    Parameters:
+        dtype: The input element type.
+
+    Args:
+        ndarray: An NDArray.
+        idx_array: An NDArray of the indices.
+        left: Left index of the partition.
+        right: Right index of the partition.
+    """
+
+    if right > left:
+        var pivot_index = left + (right - left) // 2
+        var pivot_new_index = _argsort_partition(
+            ndarray, idx_array, left, right, pivot_index
+        )
+        argsort_inplace(ndarray, idx_array, left, pivot_new_index - 1)
+        argsort_inplace(ndarray, idx_array, pivot_new_index + 1, right)
+
+
+fn argsort[
+    dtype: DType
+](ndarray: NDArray[dtype],) raises -> NDArray[DType.index]:
+    """
+    Argsort of the NDArray using quick sort algorithm.
+
+    Parameters:
+        dtype: The input element type.
+
+    Args:
+        ndarray: An NDArray.
+
+    Returns:
+        The indices of the sorted NDArray.
+    """
+
+    var array: NDArray[dtype] = ndarray
+    var length = array.size()
+
+    var idx_array = NDArray[DType.index](length)
+    for i in range(length):
+        idx_array[i] = i
+
+    argsort_inplace(array, idx_array, 0, length - 1)
+
+    return idx_array

tests/argsort.mojo

@@ -0,0 +1,19 @@
+# Test file for
+# numojo.core.sort.argsort()
+
+import numojo as nm
+from numojo.core.ndarray import NDArray
+import time
+
+fn main() raises:
+    test[nm.f64](6)
+    test[nm.i32](12)
+    test[nm.f64](1000)
+
+fn test[dtype: DType](length: Int) raises:
+    # Initialize an ND arrays of type
+    var t0 = time.now()
+    var A = NDArray[dtype](length, random=True)
+    print(A)
+    print(nm.core.sort.argsort(A))
+    print((time.now() - t0)/1e9, "s")