CODEWITHSUNDEEP
haskellvector
tsorn
tsorn

Reputation: 3625

Creating a random permutation of 1..N with Data.Vector.Unboxed.Mutable

I want to create a list containing a random permutation of the numbers 1 through N. As I understand it, it is possible to use VUM.swap in the runST, but since I need random numbers as well I figured I might do both in the IO monad.

The code below yields:

Expected type: IO (VU.Vector Int), Actual type: IO (VU.Vector (VU.Vector a0))

for the return statement. 

import qualified Data.Vector.Unboxed         as VU
import qualified Data.Vector.Unboxed.Mutable as VUM
import           System.Random

randVector :: Int -> IO (VU.Vector Int)
randVector n = do
  vector <- VU.unsafeThaw $ VU.enumFromN 1 n
  VU.forM_ (VU.fromList [2..VUM.length vector]) $ \i -> do
    j <- randomRIO(0, i) :: IO Int
    VUM.swap vector i j
  return $ VU.unsafeFreeze vector

I'm not quite sure why the return vector is nested. Do I have to use VU.fold1M_ instead?

Upvotes: 3

Views: 117

Answers (2)

Andr&#225;s Kov&#225;cs
Andr&#225;s Kov&#225;cs

Reputation: 30103

unsafeFreeze vector already returns IO (VU.Vector Int). Just change the last line to VU.unsafeFreeze vector.

On another note, you should iterate until VUM.length vector - 1, since both [x .. y] and randomRIO use inclusive ranges. Also, you can use plain forM_ here for iteration, since you only care about side effects.

import           Control.Monad
import qualified Data.Vector.Unboxed         as VU
import qualified Data.Vector.Unboxed.Mutable as VUM
import           System.Random

randVector :: Int -> IO (VU.Vector Int)
randVector n = do
  vector <- VU.unsafeThaw $ VU.enumFromN 1 n
  forM_ [2..VUM.length vector - 1] $ \i -> do
    j <- randomRIO(0, i) :: IO Int
    VUM.swap vector i j
  VU.unsafeFreeze vector

I looked at the generated code, and it seems that with GHC 7.10.3 forM_ compiles to an efficient loop while VU.forM_ retains the intermediate list and is surely significantly slower (which was my expected outcome for forM_, but I was unsure about VU.forM_).

Upvotes: 3

ErikR
ErikR

Reputation: 52029

I would try (note update at end):

import Control.Monad

randVector :: Int -> IO (VU.Vector Int)
randVector n = do
  vector <- VU.unsafeThaw $ VU.enumFromN 1 n
  forM_ [2..VUM.length vector] $ \i -> do
    j <- randomRIO(0, i) :: IO Int
    VUM.swap vector i j
  return $ VU.unsafeFreeze vector

Edit: as @András Kovács pointed out, you don't want the return at the end so the last line should be:

  VU.unsafeFreeze vector

Upvotes: 1

Related Questions