Implement sort and argsort for one dimensional SimpleArray in C++

cpp/modmesh/buffer/SimpleArray.hpp

@@ -34,6 +34,7 @@
 #include <stdexcept>
 #include <functional>
 #include <numeric>
+#include <algorithm>
 
 #if defined(_MSC_VER)
 #include <BaseTsd.h>
@@ -192,6 +193,18 @@ class SimpleArrayMixinCalculators
         }
         return ret;
     }
+
+    void sort(void)
+    {
+        auto athis = static_cast<A *>(this);
+        if (athis->ndim() != 1)
+        {
+            throw std::runtime_error(Formatter() << "SimpleArray: sorting is only supported in 1D array. "
+                                                 << athis->ndim() << "D array is currently not supported");
+        }
+
+        std::sort(athis->begin(), athis->end());
+    }
 }; /* end class SimpleArrayMixinCalculators */
 
 } /* end namespace detail */
@@ -585,6 +598,33 @@ class SimpleArray
         }
     }
 
+    SimpleArray<uint64_t> argsort(void)
+    {
+        if (ndim() != 1)
+        {
+            throw std::runtime_error(Formatter() << "SimpleArray: sorting is only supported in 1D array. "
+                                                 << ndim() << "D array is currently not supported");
+        }
+
+        SimpleArray<uint64_t> ret(shape());
+
+        { // Return array initialization
+            uint64_t cnt = 0;
+            std::for_each(ret.begin(), ret.end(), [&cnt](uint64_t & v)
+                          { v = cnt++; });
+        }
+
+        value_type const * buf = body();
+        auto cmp = [buf](uint64_t a, uint64_t b)
+        {
+            return buf[a] < buf[b];
+        };
+        std::sort(ret.begin(), ret.end(), cmp);
+        return ret;
+    }
+
+    void take_along_axis(SimpleArray<uint64_t> const & indices);
+
     template <typename... Args>
     value_type const & operator()(Args... args) const { return *vptr(args...); }
     template <typename... Args>
@@ -729,6 +769,64 @@ class SimpleArray
     value_type * m_body = nullptr;
 }; /* end class SimpleArray */
 
+template <typename T>
+void SimpleArray<T>::take_along_axis(SimpleArray<uint64_t> const & indices)
+{
+    if (indices.ndim() != 1)
+    {
+        throw std::runtime_error(Formatter() << "SimpleArray: sorting is only supported in 1D array. "
+                                             << indices.ndim() << "D indices is not supported.");
+    }
+    if (ndim() != 1)
+    {
+        throw std::runtime_error(Formatter() << "SimpleArray: sorting is only supported in 1D array. "
+                                             << ndim() << "D array is not supported to be sorted.");
+    }
+    if (shape()[0] != indices.shape()[0])
+    {
+        throw std::runtime_error(Formatter() << "SimpleArray: take_along_axis only support same shape of indices and array."
+                                             << "Array size " << shape()[0] << " != indices shape " << indices.shape()[0]);
+    }
+    if (shape()[0] < 2)
+    {
+        return;
+    }
+
+    std::vector<bool> applied_arg(shape()[0], false);
+
+    auto next = [](std::vector<bool> & vec, ssize_t last)
+    {
+        for (ssize_t i = last; i < static_cast<ssize_t>(vec.size()); ++i)
+        {
+            if (vec.at(i) == false)
+            {
+                return i;
+            }
+        }
+        return static_cast<ssize_t>(-1);
+    };
+
+    ssize_t idx = 0;
+    while ((idx = next(applied_arg, idx)) != -1)
+    {
+        value_type val = at(idx);
+
+        ssize_t dst_idx = idx;
+        ssize_t src_idx = indices[dst_idx];
+
+        while (src_idx != idx)
+        {
+            at(dst_idx) = at(src_idx);
+            applied_arg.at(dst_idx) = true;
+            dst_idx = src_idx;
+            src_idx = indices[dst_idx];
+        }
+
+        at(dst_idx) = val;
+        applied_arg.at(dst_idx) = true;
+    }
+}
+
 template <typename S>
 using is_simple_array = std::is_same<
     std::remove_reference_t<S>,

cpp/modmesh/buffer/pymod/wrap_SimpleArray.cpp

@@ -147,6 +147,14 @@ class MODMESH_PYTHON_WRAPPER_VISIBILITY WrapSimpleArray
                 "reshape",
                 [](wrapped_type const & self, py::object const & shape)
                 { return self.reshape(make_shape(shape)); })
+            .def(
+                "sort",
+                [](wrapped_type & self)
+                { self.sort(); })
+            .def(
+                "argsort",
+                [](wrapped_type & self)
+                { return pybind11::cast(self.argsort()); })
             .def_property_readonly("has_ghost", &wrapped_type::has_ghost)
             .def_property("nghost", &wrapped_type::nghost, &wrapped_type::set_nghost)
             .def_property_readonly("nbody", &wrapped_type::nbody)

tests/test_buffer.py

@@ -824,6 +824,17 @@ def test_SimpleArray_SimpleArrayPlex_type_switch(self):
         self.assertEqual(
             str(type(arrayplex_int32_2)), "<class '_modmesh.SimpleArray'>")
 
+    def test_sort(self):
+        narr = np.random.randint(0, 100, 20, dtype='int32')
+        sarr = modmesh.SimpleArrayInt32(array=narr)
+        args = sarr.argsort()
+        for i in range(1, len(args)):
+            self.assertLessEqual(sarr[args[i]], sarr[args[i]])
+
+        sarr.sort()
+        for i in range(1, len(args)):
+            self.assertLessEqual(sarr[i - 1], sarr[i])
+
 
 class SimpleArrayCalculatorsTC(unittest.TestCase):