-- | Miscellaneous vector methods.
--
-- @since 1.2.2.0
module AtCoder.Extra.Vector
  ( argsort,
  )
where

import Data.Vector.Algorithms.Intro qualified as VAI
import Data.Vector.Generic qualified as VG
import Data.Vector.Unboxed qualified as VU

-- | \(O(n \log n)\) Returns indices of the vector, stably sorted by their value.
--
-- ==== Example
-- >>> import Data.Vector.Algorithms.Intro qualified as VAI
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> argsort $ VU.fromList [0, 1, 0, 1, 0]
-- [0,2,4,1,3]
--
-- @since 1.2.3.0
{-# INLINEABLE argsort #-}
argsort :: (Ord a, VU.Unbox a) => VU.Vector a -> VU.Vector Int
argsort :: forall a. (Ord a, Unbox a) => Vector a -> Vector Int
argsort Vector a
xs =
  (forall s. MVector s Int -> ST s ()) -> Vector Int -> Vector Int
forall a.
Unbox a =>
(forall s. MVector s a -> ST s ()) -> Vector a -> Vector a
VU.modify
    ( Comparison Int -> MVector (PrimState (ST s)) Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> m ()
VAI.sortBy
        ( \Int
i Int
j ->
            ( a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Vector a
xs Vector a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i) (Vector a
xs Vector a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
j) Ordering -> Ordering -> Ordering
forall a. Semigroup a => a -> a -> a
<> Comparison Int
forall a. Ord a => a -> a -> Ordering
compare Int
i Int
j
            )
        )
    )
    (Vector Int -> Vector Int) -> Vector Int -> Vector Int
forall a b. (a -> b) -> a -> b
$ Int -> (Int -> Int) -> Vector Int
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
xs) Int -> Int
forall a. a -> a
id

-- TODO: maybe add lexical permutations, combinations and subsequences.