Leave at most max_dup consecutive duplicate elements in a list

Question

I need to remove excessive duplicates of a particular element from the list, that is, if there's more than max_dup duplicates of selected elements in a slice of the list, leave at most max_dup elements in place. Obviously, stability (retaining order of elements in the list) is a must.

Example:

>>> remove_excessive_duplicates([1, '', '', '', 2, '', 3, '', '', '', '', '', 4, ''])
[1, '', '', 2, '', 3, '', '', 4, '']

The code I wrote:

def remove_excessive_duplicates(lst, dup='', max_dup=2):
    new_lst = []
    dcnt = 0
    for elem in lst:
        if elem == dup:
            dcnt += 1 
        else:
            dcnt = 0
        if dcnt > max_dup:
            continue
        new_lst.append(elem)
    return new_lst

It works, but I'm wondering if there's some cleverer way to do it, or the one that uses logic of some Python builtins to achieve the same goal.

Answer 1

This is a one-liner using the features of the itertools module. I'll try to show how you might find this via a step-by-step transformation.

1. When a function returns a list of items, it's often more convenient (in Python) to refactor it so that it generates the items one at a time using the yield statement. This has some advantages: (i) if the caller is also processing the items one at a time, then you don't need to allocate the whole list in memory; (ii) the first item becomes available immediately, without having to wait for all the items; (iii) if the items are arriving from a file or some other form of I/O, you may be able to process early items in parallel with loading later ones; (iv) a generator can operate on a potentially infinite stream of items without running out of memory.

   def remove_excessive_duplicates(iterable, dup='', max_dup=2):
       """Generate items in iterable, but at most max_dup of any consecutive
       sequence of items equal to dup.
   
       >>> input = [1, 0, 0, 0, 2, 0, 3, 0, 0, 0, 0, 0, 4, 0]
       >>> list(remove_excessive_duplicates(input, dup=0))
       [1, 0, 0, 2, 0, 3, 0, 0, 4, 0]
   
       """
       count = 0
       for item in iterable:
           if item == dup:
               count += 1 
           else:
               count = 0
           if count <= max_dup:
               yield item
   

   (I've also taken the opportunity of adding a docstring and a doctest, and improving some of the variable names.)

2. What's going on here is that we are finding groups of identical items in the input, and for this Python provides itertools.groupby:

   for key, group in groupby(iterable):
       if key == dup:
           count = 0
           for item in group:
               if count >= max_dup:
                   break
               yield item
               count += 1
       else:
           yield from group
   

   This has the advantage that we don't need to maintain the running count in the else case, saving some work.

3. In the if case, we want to yield at most max_dup items from the iterable group. This can be done using itertools.islice:

   for key, group in groupby(iterable):
       if key == dup:
           yield from islice(group, max_dup)
       else:
           yield from group
   

4. We can factor out the repeated yield from:

   for key, group in groupby(iterable):
       yield from islice(group, max_dup) if key == dup else group
   

5. That would a sensible place to stop. But to make this into a single expression we can note that the effect of this code is that for each group in the iterable, we construct an output iterable (either islice(...) or group, depending on the key), and the result consists of the concatenation of these iterables. Concatenating a bunch of iterables is a job for the obscurely named itertools.chain.from_iterable:

   return chain.from_iterable(islice(group, max_dup) if key == dup else group
                              for key, group in groupby(iterable))