How can I find out the size of the group a given element exists in?

Question

If I have a list like such:

lst = [4,4,4,5,3,3,9]

I can use group to create this: [[4,4,4],[5],[3,3],[9]].

If I have an index pointing into the list, e.g., i = 4, which points at the first 3 element, how can I best write a function taking a list and index, and returning the size of the group to which the index points (in this case, 2, the length of the group of 3s)?

I know I can write it pretty imperatively, and that I could also probably knock something together with map length . group and count how many groups I pass through when locating my index, but it seems to me that there must be a nicer way to do this. All suggestions welcome.

Answer 1

Here's a solution in rather explicit stages so I can explain each step.

Make the grouped list:

lst = [4,4,4,5,3,3,9]
lst1 = group lst

Next, replace the groups with lengths:

lst2 = map length lst1

This gives us lst2 = [3,1,2,1]

Now replace the lengths with lists of that length that also contain the length in each position:

lst3 = map (\l -> replicate l l) lst2

Now we have lst3 = [[3,3,3],[1],[2,2],[1]], which is nearly what we need but we don't want the inner structure:

lst4 = concat lst3

Now we have something the same length as the original list, but with group lengths in each position instead of the original values. So lst4!!4 = 2.

This solution is a simplification of my original solution using an idea from the comments.

As also pointed out in that comment, a further couple of tweaks are possible - elide the two maps and the concat as concatMap, and use join in the (->) r instance of Monad to "simplify" \l -> replicate l l, giving this for the whole pipeline:

concatMap (join replicate . length) . group $ lst

Answer 2

If you have a function to compute the cumulative sum, then you can make it fairly straightforward:

cumsum :: Num a => [a] -> [a]
cumsum xs = go xs 0
    where
        go [] _ = []
        go (x:xs) i = let n = x + i in n : go xs n

groupSizeAt :: [[a]] -> Int -> Int
groupSizeAt xs i =
    let groups = zip xs $ cumsum $ map length xs
        targetGroup = head $ dropWhile ((<= i) . snd) groups
    in length $ fst targetGroup

The logic behind this is to compute the cumulative sum of the lengths of the groups, then drop groups until this cumulative sum is greater than the index you're searching for, then grab the first one, drop the cumulative sum, and compute the length again. There's a bit of redundancy, but nothing that would be too expensive for small groups (if your group size is on the order of 100s or 1000s, you might want to modify it to keep from having to compute the length of the target group twice)