{- |
Description: Group into disjoint subsets
Copyright: Copyright (C) 2023 Yoo Chung
License: GPL-3.0-or-later
Maintainer: dev@chungyc.org

Some solutions to "Problems.P27" of Ninety-Nine Haskell "Problems".
-}
module Solutions.P27 (disjointGroups) where

{- |
Given a list with the sizes of each group and list of items,
write a function to return the list of disjoint groups.
-}
disjointGroups :: [Int] -> [a] -> [[[a]]]
disjointGroups :: forall a. [Int] -> [a] -> [[[a]]]
disjointGroups [] []     = [[]]
disjointGroups [] [a]
xs     = [[[a]
xs]]  -- implicitly another group if group sizes do not sum to length
disjointGroups (Int
n:[Int]
ns) [a]
xs = (([a], [a]) -> [[[a]]]) -> [([a], [a])] -> [[[a]]]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\([a]
g,[a]
xs') -> ([[a]] -> [[a]]) -> [[[a]]] -> [[[a]]]
forall a b. (a -> b) -> [a] -> [b]
map ([a]
g:) ([[[a]]] -> [[[a]]]) -> [[[a]]] -> [[[a]]]
forall a b. (a -> b) -> a -> b
$ [Int] -> [a] -> [[[a]]]
forall a. [Int] -> [a] -> [[[a]]]
disjointGroups [Int]
ns [a]
xs') ([([a], [a])] -> [[[a]]]) -> [([a], [a])] -> [[[a]]]
forall a b. (a -> b) -> a -> b
$ Int -> [a] -> [([a], [a])]
forall a. Int -> [a] -> [([a], [a])]
extract Int
n [a]
xs

extract :: Int -> [a] -> [([a], [a])]
extract :: forall a. Int -> [a] -> [([a], [a])]
extract Int
n [a]
xs
  | [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
n = [([a]
xs,[])]  -- last group may be smaller than given size
  | Bool
otherwise      = Int -> ([a], [a]) -> [([a], [a])]
forall a. Int -> ([a], [a]) -> [([a], [a])]
extract' Int
n ([],[a]
xs)

extract' :: Int -> ([a], [a]) -> [([a],[a])]
extract' :: forall a. Int -> ([a], [a]) -> [([a], [a])]
extract' Int
0 ([a], [a])
r        = [([a], [a])
r]
extract' Int
_ ([a]
_, [])  = []
extract' Int
n ([a]
g,a
x:[a]
xs) = (([a], [a]) -> ([a], [a])) -> [([a], [a])] -> [([a], [a])]
forall a b. (a -> b) -> [a] -> [b]
map (\([a]
g',[a]
xs') -> ([a]
g',a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
xs')) (Int -> ([a], [a]) -> [([a], [a])]
forall a. Int -> ([a], [a]) -> [([a], [a])]
extract' Int
n ([a]
g,[a]
xs)) [([a], [a])] -> [([a], [a])] -> [([a], [a])]
forall a. [a] -> [a] -> [a]
++ Int -> ([a], [a]) -> [([a], [a])]
forall a. Int -> ([a], [a]) -> [([a], [a])]
extract' (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
g,[a]
xs)