Given list a = [1, 2, 2, 3] and its sublist b = [1, 2] find a list complementing b in such a way that sorted(a) == sorted(b + complement). In the example above the complement would be a list of [2, 3].
It is tempting to use list comprehension:
complement = [x for x in a if x not in b]
or sets:
complement = list(set(a) - set(b))
However, both of this ways will return complement = [3].
An obvious way of doing it would be:
complement = a[:]
for element in b:
complement.remove(element)
But that feels deeply unsatisfying and not very Pythonic. Am I missing an obvious idiom or is this the way?
As pointed out below what about performance this is O(n^2) Is there more efficient way?
The only more declarative and thus Pythonic way that pops into my mind and that improves performance for large b (and a) is to use some sort of counter with decrement:
from collections import Counter
class DecrementCounter(Counter):
def decrement(self,x):
if self[x]:
self[x] -= 1
return True
return False
Now we can use list comprehension:
b_count = DecrementCounter(b)
complement = [x for x in a if not b_count.decrement(x)]
Here we thus keep track of the counts in b, for each element in a we look whether it is part of b_count. If that is indeed the case we decrement the counter and ignore the element. Otherwise we add it to the complement. Note that this only works, if we are sure such complement exists.
After you have constructed the complement, you can check if the complement exists with:
not bool(+b_count)
If this is False, then such complement cannot be constructed (for instance a=[1] and b=[1,3]). So a full implementation could be:
b_count = DecrementCounter(b)
complement = [x for x in a if not b_count.decrement(x)]
if +b_count:
raise ValueError('complement cannot be constructed')
If dictionary lookup runs in O(1) (which it usually does, only in rare occasions it is O(n)), then this algorithm runs in O(|a|+|b|) (so the sum of the sizes of the lists). Whereas the remove approach will usually run in O(|a|×|b|).
In order to reduce complexity to your already valid approach, you could use collections.Counter (which is a specialized dictionary with fast lookup) to count items in both lists.
Then update the count by substracting values, and in the end filter the list by only keeping items whose count is > 0 and rebuild it/chain it using itertools.chain
from collections import Counter
import itertools
a = [1, 2, 2, 2, 3]
b = [1, 2]
print(list(itertools.chain.from_iterable(x*[k] for k,x in (Counter(a)-Counter(b)).items() if x > 0)))
result:
[2, 2, 3]
O(n log n)
a = [1, 2, 2, 3]
b = [1, 2]
a.sort()
b.sort()
L = []
i = j = 0
while i < len(a) and j < len(b):
if a[i] < b[j]:
L.append(a[i])
i += 1
elif a[i] > b[j]:
L.append(b[j])
j += 1
else:
i += 1
j += 1
while i < len(a):
L.append(a[i])
i += 1
while j < len(b):
L.append(b[j])
j += 1
print(L)