Unique elements of a list and corresponding indices in Haskell

Question

Given a list vs, I want to get the list vs' of the unique elements of vs, as well as the indices of the elements of vs in vs'. For example, if vs = [7, 8, 7, 8, 9] I want to get [7,8,9] (the unique elements) and [0,1,0,1,2] (the indices).

An easy implementation is:

import           Data.List
import           Data.Maybe

unique :: Eq a => [a] -> ([a], [Int])
unique vs = (vsnub, indices)
  where
  vsnub = nub vs
  indices = map (\v -> fromJust $ elemIndex v vsnub) vs

However I think this is not efficient. I've done an implementation using mutable vectors. The pseudo-code is (saying that vs is a list of numbers):

n = length of vs
idx = list of n integers (to store the indices)
visited = [false, false, ..., false] (n elements)
nvs = list of n numbers (to store the unique elements)
count = 0 
for(i = 0; i < n; ++i)
{
  if(not visited[i])
  {
    nvs[count] = vs[i]
    idx[i] = count
    visited[i] = true
    for(j = i+1; j < n; ++j)
    {
      if(vs[j] = vs[i])
      {
        visited[j] = true
        idx[j] = count
      }
    }
    count ++
  }
}
nvs = first 'count' elements of nvs

And here is my Haskell code:

{-# LANGUAGE ScopedTypeVariables #-}
import           Control.Monad               ((>>=))
import           Data.Vector.Unboxed         (Unbox, Vector, freeze, (!))
import           Data.Vector.Unboxed.Mutable (IOVector, new, write)
import qualified Data.Vector.Unboxed.Mutable as VM

unique' :: forall a . (Unbox a, Eq a) => [a] -> IO (Vector a, Vector Int)
unique' vs = do
  let n = length vs
  idx <- VM.replicate n 0 :: IO (IOVector Int)
  visited <- VM.replicate n False :: IO (IOVector Bool)
  nvs <- new n :: IO (IOVector a)
  let inner :: Int -> Int -> Int -> IO ()
      inner i j count | j == n = return ()
                      | otherwise =
                        if vs !! i == vs !! j
                          then do
                            write visited j True
                            write idx j count
                            inner i (j+1) count
                          else inner i (j+1) count
  let go :: Int -> Int -> IO (IOVector a)
      go i count | i == n = return $ VM.take count nvs
                 | otherwise = do
                   vst <- VM.read visited i
                   if not vst
                     then do
                       write nvs count (vs !! i)
                       write idx i count
                       write visited i True
                       _ <- inner i (i+1) count
                       go (i+1) (count + 1)
                     else go (i+1) count
  nvs' <- go 0 0 >>= freeze
  idx' <- freeze idx
  return (nvs', idx')

Is it nice? Can it be improved? Is there a solution that does not resort to IO?

Answer 1

Various Comments

First, Stack Exchange sent the body of my response into the void when I logged in to post it, so I'm going to be a bit more brief and a bit less organized than I originally wanted because I'm lazy.

This is not really reviewing your code, so it would've been better served as a comment, but I didn't think to comment my approach before I wrote out all the code. Might as well share it in case it's helpful.

To be honest, I don't quite understand your second solution, but on a second read it looks like it might still be n^2 (same asymptotic complexity as your first solution), so that would explain why it's slower. I assumed at first it was just something to do with mutability and vectors (which I don't really understand myself), but on second read I'm no longer so sure about this.

If you don't mind me being frank in answering your first question: I don't really think it's nice. I had a pretty rough time reading it, to the point where I instead just wrote a more Haskell-inclined solution that has better asymptotic performance than your first one.

Proposed Solution

Are you okay with an n log n solution? If memory serves, with the right choice of set/map, the log's base is so large that it is effectively linear.

Here is one such solution

import Data.List (foldl', sortOn)
import Data.Map (Map)
import qualified Data.Map as Map

-- | Gives the unique elements of 'elems' in order.
orderedNub :: Ord a => [a] -> [a]
orderedNub elems = sortOn (firstIndexMap Map.!) (Map.keys firstIndexMap)
  where
    -- Insert such that if the value is already in 'firstIndexMap', it is not
    -- updated.
    addElem m (elem, index) = Map.insertWith (flip const) elem index m
    -- Left fold in order to get the first occurance of each element.
    firstIndexMap = foldl' addElem Map.empty $ zip elems [0..]

-- | Gives the unique elements of 'elems' and the indices of 'elems' in
-- the unique list of elements.
unique :: Ord a => [a] -> ([a], [Int])
unique elems = (uniques, indices)
  where
    uniques = orderedNub elems
    uniqueInds = zip uniques [0..]
    indexMap = foldr (uncurry Map.insert) Map.empty uniqueInds
    -- We can use unsafe indexing since we know that 'indexMap' has
    -- the right values.
    indices = [indexMap Map.! x | x <- elems]

Some Notes About the Solution

1. I changed the code so it is correct now, assuming you want to preserve order. Some of these comments may be a bit out of date as a result.
2. orderedNub' is a kind of clunky definition, but it should have better asymptotic performance than nub. I’m pretty sure you can avoid the overhead of an actual sort (one clunky way would be to use a vector like you’re using).
3. uncurry Map.insert is a cute pointfree thing, if that's not your dig just sub it with a lambda or helper function.
4. I can't think of a good way to avoid using Map.!, even though it's partial. If I used Map.lookup, then there'd have to be a fromJust or a Maybe overhead, even though the map is, by construction, going to have the right keys. Guess this is a shortcoming of the type system.