diff --git a/src/Data/Set.purs b/src/Data/Set.purs index 0269cd5..022e390 100644 --- a/src/Data/Set.purs +++ b/src/Data/Set.purs @@ -24,6 +24,7 @@ module Data.Set , subset , properSubset , intersection + , choose ) where import Prelude hiding (map) @@ -176,3 +177,11 @@ intersection s1 s2 = fromFoldable $ runPure (runSTArray (emptySTArray >>= inters LT -> pure $ Loop {a: l + 1, b: r} GT -> pure $ Loop {a: l, b: r + 1} else pure $ Done acc + +-- | The set of all subsets of specified size +choose :: forall a. Ord a => Set a -> Int -> Set (Set a) +choose = go (singleton empty) + where + go acc _ 0 = acc + go acc s n = + unions $ map (\e -> go (map (insert e) acc) (delete e s) (n - 1)) s diff --git a/test/Test/Main.purs b/test/Test/Main.purs index 66c9e74..be08033 100644 --- a/test/Test/Main.purs +++ b/test/Test/Main.purs @@ -26,3 +26,11 @@ main = do s2 = S.fromFoldable [2,4,6,8,10] s3 = S.fromFoldable [2,4] assert $ S.intersection s1 s2 == s3 + + log "choose" + do let set = S.fromFoldable [0, 1, 2] + assert $ S.choose set 2 == S.fromFoldable `S.map` S.fromFoldable [ + [0, 1] + , [1, 2] + , [0, 2] + ]