高效的Haskell相当于NumPy的argsort

Question

是否有一个标准的Haskell等同于NumPy的argsort功能？

我正在使用HMatrix,因此,希望兼容的功能Vector R是别名Data.Vector.Storable.Vector Double.argSort下面的函数是我目前使用的实现:

{-# 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]

我正在使用显式限定的imports来明确每个类型和函数的来源.

此实现不是非常有效,因为它将输入向量转换为列表并将结果转换回向量.在某个地方存在这样的事情(但效率更高)？

更新

@leftaroundabout有一个很好的解决方案.这是我最终得到的解决方案:

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由于数据向量是a,因此更直接可用Storable.这为索引使用了未装箱的矢量.

Answer 1

使用矢量算法:

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来往于可存储的载体.