How to generalize picking subsets of size `n` from `m` alternatives with state between picks in Haskell

Question

I'm working on brute-forcing solutions to a maths puzzle, and am struggling to abstract my solution so that I can easily solve varying puzzle sizes.

The puzzle can be found at https://www.think-maths.co.uk/uniquedistance. Please don't read on if you want to solve the puzzle yourself without spoilers. If you just want to help me solve the Haskell programming problem at hand, you don't need to spend time on studying what the puzzle is.

What I'm trying to do in my solution shown below, is to find subsets of size n from a pool of n^2 different options, such that some binary function metric produces unique results for all pairs of options from the chosen subset.

At first I wrote a solution along the lines of

combinations :: Int -> [a] -> [[a]]
combinations 0 _ = [[]]
combinations _ [] = []
combinations n xs = [ a:rec | (a:as) <- tails xs, rec <- combinations (pred n) as ]

which gave me all the possible subsets, and subsequently checked whether or not any individual subset satisfied the uniqueness requirement of the given metric for all possible pairs chosen from the subset using

import qualified Data.IntSet as IS

check :: [a] -> Bool
check = noDupes . metrics
  where metrics ps = [ metric a b | (a:bs) <- tails ps, b <- bs ]
        noDupes = go IS.empty
        go _ [] = True
        go s (x:xs) | IS.member x s = False
                    | otherwise = go (IS.insert x s) xs

From there, filter check (combinations n) would give me correct solutions for any given n. However, to improve performance, I wanted to change my computation such that rather than first generating all subsets of size n and only then checking if my constraint holds for the full subset, it would instead discard subsets smaller than n elements earlier on, allowing me to compute the expensive metric less often.

I wasn't easily able to transform my above solution into what I wanted, but so far I've been able to come up with the following (which also includes some more concrete types and a definition of the metric, but I think you can ignore that if you don't care about the details of the puzzle):

import qualified Data.IntSet as IS
import Data.Maybe
import Control.Monad
import Data.List
import Linear.V2 (V2(..))

-- euclidean distance squared
metric :: V2 Int -> V2 Int -> Int
metric (V2 x1 y1) (V2 x2 y2) = ((x1-x2)^2) + ((y1-y2)^2)

-- metric of a new candidate point to all previous points
metrics p = map (metric p)

-- check if the previously seen set of metrics are compatible with the metrics
-- of a new candidate. Nothing if they're not, and Just the union of the
-- previous and new metrics.
checkCompatibility :: IS.IntSet -> [Int] -> Maybe IS.IntSet
checkCompatibility s [] = Just s
checkCompatibility s (x:xs) | IS.member x s = Nothing
                            | otherwise = checkCompatibility (IS.insert x s) xs

-- all combinations of choosing 1 points from the input
combinations1 :: [V2 Int] -> [[V2 Int]]
combinations1 xs = do
  (a:bs) <- tails xs
  let ret = [a]

  return ret

-- all combinations of choosing 2 points from the input
combinations2 :: [V2 Int] -> [[V2 Int]]
combinations2 xs = do
  (a:bs) <- tails xs
  let ret = [a]

  (b:cs) <- tails bs
  let sset = checkCompatibility IS.empty (metrics b ret)
  guard (maybe False (not . IS.null) sset)
  let ret' = b:ret

  return (reverse ret')

-- all combinations of choosing 3 points from the input, where the "metric" between any pair of points is unique
combinations3 :: [V2 Int] -> [[V2 Int]]
combinations3 xs = do
  (a:bs) <- tails xs
  let ret = [a]

  (b:cs) <- tails bs
  let sset = checkCompatibility IS.empty (metrics b ret)
  guard (maybe False (not . IS.null) sset)
  let ret' = b:ret

  (c:ds) <- tails cs
  let sset' = checkCompatibility (fromJust sset) (metrics c ret')
  guard (maybe False (not . IS.null) sset')
  let ret'' = c:ret'

  return (reverse ret'')

-- all combinations of choosing 4 points from the input, where the "metric" between any pair of points is unique
combinations4 :: [V2 Int] -> [[V2 Int]]
combinations4 xs = do
  (a:bs) <- tails xs
  let ret = [a]

  (b:cs) <- tails bs
  let sset = checkCompatibility IS.empty (metrics b ret)
  guard (maybe False (not . IS.null) sset)
  let ret' = b:ret

  (c:ds) <- tails cs
  let sset' = checkCompatibility (fromJust sset) (metrics c ret')
  guard (maybe False (not . IS.null) sset')
  let ret'' = c:ret'

  (d:es) <- tails ds
  let sset'' = checkCompatibility (fromJust sset') (metrics d ret'')
  guard (maybe False (not . IS.null) sset'')
  let ret''' = d:ret''

  return (reverse ret''')

combinations5 :: [V2 Int] -> [[V2 Int]]
combinations5 xs = do
  (a:bs) <- tails xs
  let ret = [a]

  (b:cs) <- tails bs
  let sset = checkCompatibility IS.empty (metrics b ret)
  guard (maybe False (not . IS.null) sset)
  let ret' = b:ret

  (c:ds) <- tails cs
  let sset' = checkCompatibility (fromJust sset) (metrics c ret')
  guard (maybe False (not . IS.null) sset')
  let ret'' = c:ret'

  (d:es) <- tails ds
  let sset'' = checkCompatibility (fromJust sset') (metrics d ret'')
  guard (maybe False (not . IS.null) sset'')
  let ret''' = d:ret''

  (e:fs) <- tails es
  let sset''' = checkCompatibility (fromJust sset'') (metrics e ret''')
  guard (maybe False (not . IS.null) sset''')
  let ret'''' = e:ret'''

  return (reverse ret'''')

combinations6 :: [V2 Int] -> [[V2 Int]]
combinations6 xs = do
  (a:bs) <- tails xs
  let ret = [a]

  (b:cs) <- tails bs
  let sset = checkCompatibility IS.empty (metrics b ret)
  guard (maybe False (not . IS.null) sset)
  let ret' = b:ret

  (c:ds) <- tails cs
  let sset' = checkCompatibility (fromJust sset) (metrics c ret')
  guard (maybe False (not . IS.null) sset')
  let ret'' = c:ret'

  (d:es) <- tails ds
  let sset'' = checkCompatibility (fromJust sset') (metrics d ret'')
  guard (maybe False (not . IS.null) sset'')
  let ret''' = d:ret''

  (e:fs) <- tails es
  let sset''' = checkCompatibility (fromJust sset'') (metrics e ret''')
  guard (maybe False (not . IS.null) sset''')
  let ret'''' = e:ret'''

  (f:gs) <- tails fs
  let sset'''' = checkCompatibility (fromJust sset''') (metrics f ret'''')
  guard (maybe False (not . IS.null) sset'''')
  let ret''''' = f:ret''''

  return (reverse ret''''')

bruteforce :: Int -> ([V2 Int] -> [[V2 Int]]) -> [[V2 Int]]
bruteforce n f = f positions
  where positions = [ V2 x y | x <- [0..pred n], y <- [0..pred n] ]

Do note how the various implementations for different values of n are extremely similar in the same way that my original combinations function from above would've been if I hadn't written it recursively with an n parameter.

What I'm trying to work out is how to parameterize my combinations1, combinations2, combinations3, and so on functions such that I don't have to tediously write a solution for every value of n.

-- all combinations of choosing n points from the input, where the "metric" between any pair of points is unique
combinationsN :: Int -> [V2 Int] -> [[V2 Int]]
combinationsN 0 _ = [[]]
combinationsN _ [] = []
combinationsN n xs = undefined

For educational purposes, I think I'd mainly be interested in how to accomplish this while manually stringing the state along between steps so I could later refine that to a solution using Control.Monad.State, but I'd also be interested in seeing other approaches to maintain the state between the steps.

I'd also appreciate suggestions for a better question title. Not really knowing how to do what I want, I really don't know what terms I should be using to ask about it.

Thanks!

Answer 1

Well, you have the idea. Grow the IntSet through combinations. You can do this by adding extra parameters to combinations:

solve :: Int -> [V2 Int] -> [[V2 Int]]
solve n xs = go n xs IS.empty []
  where go :: Int -> [V2 Int] -> IntSet -> [V2 Int] -> [[V2 Int]]
        go 0 _  _       seen = [reverse seen]
        go n xs metrics seen = [ rec
                               | (a : as) <- tails xs
                               , metrics' <- maybeToList $ addMetrics a seen metrics
                               , rec <- go (pred n) as metrics' (a : seen)]
        addMetrics :: V2 Int -> [V2 Int] -> IntSet -> Maybe IntSet
        addMetrics _ [] i = Just i
        addMetrics a (b : bs) i = do
           i' <- addMetrics a bs i
           let m = metric a b
           guard $ m `IS.notMember` i'
           return $ IS.insert m i'

This transform is very common: you keep some extra data around in some internal, probably recursive function, then throw it away when you're done. Now, to generalize: V2 Int can become a, IntSet can become s, addMetrics and IS.empty become parameters, and Maybe generalizes to [].

-- realized that it's not really "pruning" if we're allowing [s] instead of just Maybe s, but meh
pruningCombs :: s -> (a -> [a] -> s -> [s]) -> Int -> [a] -> [[a]]
pruningCombs e grow n xs = go n xs e []
    where go 0 _  _ seen = [reverse seen]
          go n xs s seen = [ rec
                           | (a : as) <- tails xs
                           , s' <- grow a seen s
                           , rec <- go (pred n) as s' (a : seen)]

solve = pruningCombs IS.empty \a -> execStateT . traverse \b -> do
    let m = metric a b
    guard =<< gets (IS.notMember m)
    modify $ IS.insert m

This runs in about the same time as your combinations6.