groupBy function which groups non-adjacent elements

Question

The groupBy function groups the elements of a list when they are equal according to a user-defined predicate and when they are adjacent in the list. Is there a function for grouping elements without taking care of the adjacency?

I deal with polynomials of three variables. Here is an example of such a polynomial written as a sum of monomials:

p = (M (Monomial {coefficient = 1.0, powers = (1,0,0)}) :+: M (Monomial {coefficient = 1.0, powers = (0,1,0)})) :+: M (Monomial {coefficient = 1.0, powers = (1,0,0)})

That means that p = 1.0 * x^1y^0z^0 + 1.0 * x^0y^1z^0 + 1.0 * x^1y^0z^0. The first term and the third term are summable: they have the same powers and then we can add them and get 2.0 * x^1y^0z^0. My goal is to do this addition, in order to simplify p.

I have a function which converts such a p to the list of the summed monomials:

>> toListOfMonomials p
[M (Monomial {coefficient = 1.0, powers = (1,0,0)}),M (Monomial {coefficient = 1.0, powers = (0,1,0)}),M (Monomial {coefficient = 1.0, powers = (1,0,0)})]

I want to group the monomials which have the same powers. If I do

groupBy ((==) `on` getPowers)
        (toListOfMonomials p)

then the two M (Monomial {coefficient = 1.0, powers = (1,0,0)}) are not grouped because they are not adjacent.

A solution consists in sorting the list according to the powers before using groupBy. In this way, two (or more) monomials having the same powers are adjacent in the sorted list. To define an order on the powers (powers is a triplet of integers), I firstly define a Cantor pairing function:

cantorPairing :: (Int, Int) -> Int
cantorPairing (k1,k2) = (k1+k2)*(k1+k2+1) + 2*k2

cantorPairing' :: (Int, Int, Int) -> Int
cantorPairing' (k1,k2,k3) = cantorPairing(cantorPairing(k1,k2),k3)

then I can compare two powers by comparing their values under the Cantor pairing function:

groupBy ((==) `on` getPowers)
        (sortBy (compare `on` (cantorPairing' . getPowers)) (toListOfMonomials p))

This gives the desired result: with this result I can easily sum the two monomials having the same powers.

But that looks heavy, doesn't it? Isn't there a groupBy function which groups also non-adjacent elements? Otherwise, is there a more rapid way to achieve the desired result?

Answer 1

For now I cannot imagine general groupBy function that would work faster than O(n^2) time, but you may use something like this

groupBy2 :: (a -> a -> Bool) -> [a] -> [[a]]
groupBy2 = go [] where
  go acc comp [] = acc
  go acc comp (h:t) =
    let (hs, nohs) = partition (comp h) t
    in go ((h:hs):acc) comp nohs

It works exactly like regular groupBy, but it joins non-adjacent element classes.

However, if you are going to use on function the problem becomes a bit simplier, as we may use it's result as key for a map:

import qualified Data.Map as M

groupOn :: (Ord b) => (a -> b) -> [a] -> [[a]]
groupOn f =
  let unpack = fmap snd . M.toList
      fld m a = case M.lookup (f a) m of
        Nothing -> M.insert (f a) [a] m
        Just as -> M.insert (f a) (a:as) m
  in unpack . foldl fld M.empty

This is more efficient equivalent of

groupOn' f = groupBy2 ((==) `on` f)

(modulo ordering)

And btw – triplets and pairs already have defined Ord instance, you may compare them just like Ints