in-place calculation of permutations

Author: nanonaren

Data/Vector/Generic/Mutable.hs

@@ -44,6 +44,7 @@ module Data.Vector.Generic.Mutable (
   unsafeRead, unsafeWrite, unsafeModify, unsafeSwap, unsafeExchange,
 
   -- * Modifying vectors
+  nextPermutation,
 
   -- ** Filling and copying
   set, copy, move, unsafeCopy, unsafeMove,
@@ -986,3 +987,38 @@ partitionUnknown f s
                       v2' <- unsafeAppend1 v2 i2 x
                       return (v1, i1, v2', i2+1)
 
+{-
+http://en.wikipedia.org/wiki/Permutation#Algorithms_to_generate_permutations
+
+The following algorithm generates the next permutation lexicographically after
+a given permutation. It changes the given permutation in-place.
+
+1. Find the largest index k such that a[k] < a[k + 1]. If no such index exists,
+   the permutation is the last permutation.
+2. Find the largest index l greater than k such that a[k] < a[l].
+3. Swap the value of a[k] with that of a[l].
+4. Reverse the sequence from a[k + 1] up to and including the final element a[n]
+-}
+
+-- | Compute the next (lexicographically) permutation of given vector in-place.
+--   Returns False when input is the last permtuation
+nextPermutation :: (PrimMonad m,Ord e,MVector v e) => v (PrimState m) e -> m Bool
+nextPermutation v
+    | dim < 2 = return False
+    | otherwise = do
+        val <- unsafeRead v 0
+        (k,l) <- loop val (-1) 0 val 1
+        if k < 0
+         then return False
+         else unsafeSwap v k l >>
+              reverse (unsafeSlice (k+1) (dim-k-1) v) >>
+              return True
+    where loop !kval !k !l !prev !i
+              | i == dim = return (k,l)
+              | otherwise  = do
+                  cur <- unsafeRead v i
+                  -- TODO: make tuple unboxed
+                  let (kval',k') = if prev < cur then (prev,i-1) else (kval,k)
+                      l' = if kval' < cur then i else l
+                  loop kval' k' l' cur (i+1)
+          dim = length v

Data/Vector/Mutable.hs

@@ -44,6 +44,7 @@ module Data.Vector.Mutable (
   unsafeRead, unsafeWrite, unsafeModify, unsafeSwap,
 
   -- * Modifying vectors
+  nextPermutation,
 
   -- ** Filling and copying
   set, copy, move, unsafeCopy, unsafeMove
@@ -404,3 +405,8 @@ unsafeMove :: PrimMonad m => MVector (PrimState m) a   -- ^ target
 {-# INLINE unsafeMove #-}
 unsafeMove = G.unsafeMove
 
+-- | Compute the next (lexicographically) permutation of given vector in-place.
+--   Returns False when input is the last permtuation
+nextPermutation :: (PrimMonad m,Ord e) => MVector (PrimState m) e -> m Bool
+{-# INLINE nextPermutation #-}
+nextPermutation = G.nextPermutation

Data/Vector/Primitive/Mutable.hs

@@ -44,6 +44,7 @@ module Data.Vector.Primitive.Mutable (
   unsafeRead, unsafeWrite, unsafeModify, unsafeSwap,
 
   -- * Modifying vectors
+  nextPermutation,
 
   -- ** Filling and copying
   set, copy, move, unsafeCopy, unsafeMove
@@ -342,3 +343,8 @@ unsafeMove :: (PrimMonad m, Prim a)
 {-# INLINE unsafeMove #-}
 unsafeMove = G.unsafeMove
 
+-- | Compute the next (lexicographically) permutation of given vector in-place.
+--   Returns False when input is the last permtuation
+nextPermutation :: (PrimMonad m,Ord e,Prim e) => MVector (PrimState m) e -> m Bool
+{-# INLINE nextPermutation #-}
+nextPermutation = G.nextPermutation

Data/Vector/Unboxed/Mutable.hs

@@ -48,6 +48,7 @@ module Data.Vector.Unboxed.Mutable (
   unsafeRead, unsafeWrite, unsafeModify, unsafeSwap,
 
   -- * Modifying vectors
+  nextPermutation,
 
   -- ** Filling and copying
   set, copy, move, unsafeCopy, unsafeMove
@@ -294,6 +295,11 @@ unsafeMove :: (PrimMonad m, Unbox a)
 {-# INLINE unsafeMove #-}
 unsafeMove = G.unsafeMove
 
+-- | Compute the next (lexicographically) permutation of given vector in-place.
+--   Returns False when input is the last permtuation
+nextPermutation :: (PrimMonad m,Ord e,Unbox e) => MVector (PrimState m) e -> m Bool
+{-# INLINE nextPermutation #-}
+nextPermutation = G.nextPermutation
+
 #define DEFINE_MUTABLE
 #include "unbox-tuple-instances"
-

tests/Tests/Move.hs

@@ -6,6 +6,10 @@ import Test.QuickCheck.Property (Property(..))
 
 import Utilities ()
 
+import Control.Monad (replicateM)
+import Control.Monad.ST (runST)
+import Data.List (sort,permutations)
+
 import qualified Data.Vector.Generic as G
 import qualified Data.Vector.Generic.Mutable as M
 
@@ -28,8 +32,18 @@ testMove v = G.length v > 0 ==> (MkProperty $ do
   actual <- return $  G.modify (\ mv -> M.move (M.slice dstOff len mv) (M.slice srcOff len mv)) v
   unProperty $ counterexample ("Move: " ++ show (v, dstOff, srcOff, len)) (expected == actual))
 
+checkPermutations :: Int -> Bool
+checkPermutations n = runST $ do
+    vec <- U.thaw (U.fromList [1..n])
+    res <- replicateM (product [1..n]) $ M.nextPermutation vec >> U.freeze vec >>= return . U.toList
+    return $! ([1..n] : res) == sort (permutations [1..n]) ++ [[n,n-1..1]]
+
+testPermutations :: Bool
+testPermutations = all checkPermutations [1..7]
+
 tests =
     [testProperty "Data.Vector.Mutable (Move)" (testMove :: V.Vector Int -> Property),
      testProperty "Data.Vector.Primitive.Mutable (Move)" (testMove :: P.Vector Int -> Property),
      testProperty "Data.Vector.Unboxed.Mutable (Move)" (testMove :: U.Vector Int -> Property),
-     testProperty "Data.Vector.Storable.Mutable (Move)" (testMove :: S.Vector Int -> Property)]
+     testProperty "Data.Vector.Storable.Mutable (Move)" (testMove :: S.Vector Int -> Property),
+     testProperty "Data.Vector.Generic.Mutable (nextPermutation)" testPermutations]