Effective Haskell equivalent for NumPy argument

Question

Is there a standard Haskell equivalent of the NumPy function argsort?

I am using HMatrix and, therefore, would like a function compatible with Vector R, which is an alias for Data.Vector.Storable.Vector Double. The function below argsortis the implementation that I am currently using:

{-# LANGUAGE NoImplicitPrelude #-}

module Main where

import qualified Data.List as L
import qualified Data.Vector as V
import qualified Data.Vector.Storable as VS
import           Prelude (($), Double, IO, Int, compare, print, snd)

a :: VS.Vector Double
a = VS.fromList [40.0, 20.0, 10.0, 11.0]

argSort :: VS.Vector Double -> V.Vector Int
argSort xs = V.fromList (L.map snd $ L.sortBy (\(x0, _) (x1, _) -> compare x0 x1) (L.zip (VS.toList xs) [0..]))

main :: IO ()
main = print $ argSort a -- yields [2,3,1,0]

I use an explicitly qualified importone to make it clear where all the types and functions come from.

This implementation is not very efficient, because it converts the input vector to a list, and the result returns to the vector. Does something like this (but more efficient) exist somewhere?

Update

@leftaroundabout has a good solution. This is the solution I came across:

module LAUtil.Sorting
  ( IndexVector
  , argSort
  )
  where

import           Control.Monad
import           Control.Monad.ST
import           Data.Ord
import qualified Data.Vector.Algorithms.Intro as VAI
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Unboxed.Mutable as VUM
import           Numeric.LinearAlgebra

type IndexVector = VU.Vector Int

argSort :: Vector R -> IndexVector
argSort xs = runST $ do
    let l = VS.length xs
    t0 <- VUM.new l
    forM_ [0..l - 1] $
        \i -> VUM.unsafeWrite t0 i (i, (VS.!) xs i)
    VAI.sortBy (comparing snd) t0
    t1 <- VUM.new l
    forM_ [0..l - 1] $
        \i -> VUM.unsafeRead t0 i >>= \(x, _) -> VUM.unsafeWrite t1 i x
    VU.freeze t1

Numeric.LinearAlgebra, Storable. unboxed .

Answer 1

vector-algorithms:

import Data.Ord (comparing)

import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Algorithms.Intro as VAlgo

argSort :: (Ord a, VU.Unbox a) => VU.Vector a -> VU.Vector Int
argSort xs = VU.map fst $ VU.create $ do
    xsi <- VU.thaw $ VU.indexed xs
    VAlgo.sortBy (comparing snd) xsi
    return xsi

, Unboxed, Storable . , C FFI . , , convert .

Answer 2

, DataMap, , . n = xs.

import Data.Map as M (toList, fromList, toAscList)

    out :: Int -> [Double] -> [Int]
    out n !xs = let !a=  (M.toAscList (M.fromList $! (zip xs [0..n])))
                    !res=a `seq` L.map snd a
                in res

, :

out 12 [1,2,3,4,1,2,3,4,1,2,3,4] == out 12 [1,2,3,4,1,3,2,4,1,2,3,4]