Algorithm to remove duplicates array leaving at most 2 elements

Question

My interview question was that I need to return the length of an array that removed duplicates but we can leave at most 2 duplicates.

For example, [1, 1, 1, 2, 2, 3] the new array would be [1, 1, 2, 2, 3]. So the new length would be 5. I came up with an algorithm with \$O(2n)\$ I believe.

def removeDuplicates(nums):
    if nums is None:
        return 0

    if len(nums) == 0:
        return 0

    if len(nums) == 1:
        return 1

    new_array = {}
    for num in nums:
        new_array[num] = new_array.get(num, 0) + 1

    new_length = 0
    for key in new_array:
        if new_array[key] > 2:
            new_length = new_length + 2
        else:
            new_length = new_length + new_array[key]

    return new_length

new_length = removeDuplicates([1, 1, 1, 2, 2, 3])
assert new_length == 5

Can I improve this code to be any faster?

Answer 1

I am interviewing myself these days, and I think I am developing a good sense of what interviewers are after. Regarding your code, here are a few things I would have done differently:

1. Know the language conventions! Python's PEP8 dictates lowercase_with_underscores format for function and variable names.

2. Avoid unnecessary explicit checks. Your function should work on any iterable, there is no point in checking for None and not for, e.g. an int or a float, which would also error out when you try to compute their len. If you want to be pedantic about it, do something like:

   from collections import Iterable
   ...
   assert isinstance(nums, Iterable)
   

   I would personally leave that check entirely out.

3. Your other checks are also pretty much unnecessary, as Python has very well defined rules for iteration over empty iterables (does nothing), and adding a special path for size-one iterables could easily be dismissed as premature optimization (the root of all evil!). Even if you ignore this advice, PEP8 prefers if not nums over if len(nums) == 0.

4. If you reorganize your code a little, you can get rid of the second loop entirely:

   counts = {}
   new_length = 0
   for num in nums:
       count = counts.get(num, 0)
       if count < 2:
           counts[num] = count + 1
           new_length += 1
   

If you put it altogether, your function would end up looking like:

def remove_duplicates(nums, dups_to_keep=2):
    counts = {}
    new_len = 0
    for num in nums:
        count = counts.get(num, 0)
        if count < dups_to_keep:
            counts[num] = count + 1
            new_len += 1
    return new_len

It may be marginally faster than your code, but I think the most important improvement is in compactness and readability.

Answer 2

You're just reimplementing collections.Counter.

def remove_duplicates(lst):
    return sum(min(x, 2) for x in collections.Counter(lst).values())

---

By the way, \$O(2n)\$ is equivalent \$O(n)\$. However, the bound is actually tighter: \$O(n + m)\$, where m is the number of distinct elements.

This is, however, an average bound, because it relies on a dict. If the worst case hash collisions happen, you get \$O(m(n+m))\$