Question

I need a function that does the same thing as itertools.combinations(iterable, r) in python

So far I came up with this:

{-| forward application -}
x -: f = f x
infixl 0 -:

{-| combinations 2 "ABCD" = ["AB","AC","AD","BC","BD","CD"] -}
combinations :: Ord a => Int -> [a] -> [[a]]
combinations k l = (sequence . replicate k) l -: map sort -: sort -: nub
    -: filter (\l -> (length . nub) l == length l)

Is there a more elegant and efficient solution?

Was it helpful?

Solution 2

(Based on @JoseJuan's answer)

You can also use a list comprehension to filter out those where the second character is not strictly smaller than the first:

[x| x <- mapM (const "ABCD") [1..2], head x < head (tail x) ]

OTHER TIPS

xs elements taken n by n is

mapM (const xs) [1..n]

all combinations (n = 1, 2, ...) is

allCombs xs = [1..] >>= \n -> mapM (const xs) [1..n]

if you need without repetition

filter ((n==).length.nub)

then

combinationsWRep xs n = filter ((n==).length.nub) $ mapM (const xs) [1..n]

(Based on @FrankSchmitt’s answer)

We have map (const x) [1..n] == replicate n x so we could change his answer to

[x| x <- sequence (replicate 2 "ABCD"), head x < head (tail x) ]

And while in original question, 2 was a parameter k, for this particular example would probably not want to replicate with 2 and write

[ [x1,x2] | x1 <- "ABCD", x2 <- "ABCD", x1 < x2 ]

instead.

With a parameter k things are a bit more tricky if you want to generate them without duplicates. I’d do it recursively:

f 0 _  = [[]]
f _ [] = []
f k as = [ x : xs | (x:as') <- tails as, xs <- f (k-1) as' ]

(This variant does not remove duplicates if there are already in the list as; if you worry about them, pass nub as to it)

This SO answer:

subsequences of length n from list performance

is the fastest solution to the problem that I've seen.

compositions :: Int -> [a] -> [[a]]
compositions k xs
    | k > length xs = []
    | k <= 0        = [[]]
    | otherwise     = csWithoutHead ++ csWithHead
    where   csWithoutHead = compositions k $ tail xs
            csWithHead    = [ head xs : ys | ys <- compositions (k - 1) $ tail xs ]
Licensed under: CC-BY-SA with attribution
Not affiliated with StackOverflow
scroll top