Is there a simple method, maybe a one-liner, to pop an item from Python set, and get default value if it does not exist?
It could be done in four lines of code and two hash table lookups:
def pop_default(coll: set, v, default=None):
if v not in coll:
return default
coll.remove(v)
return v
Can this function be re-written and
Function usage example:
s = {'a', 'b', 'c'}
assert pop_default(s, 'a') == 'a'
assert s == {'b', 'c'}
assert pop_default(s, 'x') == None
assert s == {'b', 'c'}
Two separate improvements were proposed by @jonrsharpe and @TimurShtatland.
def pop_default(s: set, v, default=None):
return s.remove(v) or v if v in s else default
A clear winner. Same efficiency, one line of code.
def pop_default(coll: set, v, default=None):
try:
coll.remove(v)
return v
except KeyError:
return default
This solution is ~5 times slower in @nocomment 's test but is faster if exception is not raised.
This requires only a single lookup into the set, so in rare corner cases it is somewhat more efficient. Unfortunately, more generally (which includes the OP case, as well as many realistic cases), it is less efficient. Also, it is longer by one line. Note that I renamed the method to avoid clashing with the existing pop
method.
def remove_with_default(coll: set, v, default=None):
try:
coll.remove(v)
return v
except KeyError:
return default
s = {'a', 'b', 'c'}
assert remove_with_default(s, 'a') == 'a'
assert s == {'b', 'c'}
assert remove_with_default(s, 'x') == None
assert s == {'b', 'c'}
(Thanks to no comment for the benchmarking code and to the OP for the suggestion of the corner case)
from timeit import repeat
from random import randrange
import numpy as np
def pop(coll: set, v, default=None):
if v not in coll:
return default
coll.remove(v)
return v
def remove_with_default(coll: set, v, default=None):
try:
coll.remove(v)
return v
except KeyError:
return default
set_size = 10
num_reps = 10**5
print("OP case:")
for f in [pop, remove_with_default] * 3:
s = {'a', 'b', 'c'}
t = 1e9 / num_reps * np.mean(repeat("f(s, 'a')", number=num_reps, globals=globals()))
print(f'{t:0.1e} ns ', f.__name__)
print("more realistic case:")
for f in [pop, remove_with_default] * 3:
t = 1e9 / num_reps * np.mean(repeat("s = set(range(set_size)); [f(s, randrange(set_size)) for i in range(set_size)]", number=num_reps, globals=globals()))
print(f'{t:0.1e} ns ', f.__name__)
print("corner case:")
for f in [pop, remove_with_default] * 3:
t = 1e9 / num_reps * np.mean(repeat("s = set(range(set_size)); [f(s, i) for i in range(set_size)]", number=num_reps, globals=globals()))
print(f'{t:0.1e} ns ', f.__name__)
Results:
OP case:
6.4e+01 ns pop
1.7e+02 ns remove_with_default
3.3e+01 ns pop
1.6e+02 ns remove_with_default
3.3e+01 ns pop
1.6e+02 ns remove_with_default
more realistic case:
2.8e+03 ns pop
3.4e+03 ns remove_with_default
2.8e+03 ns pop
3.3e+03 ns remove_with_default
2.8e+03 ns pop
3.4e+03 ns remove_with_default
corner case:
8.2e+02 ns pop
7.4e+02 ns remove_with_default
8.2e+02 ns pop
7.1e+02 ns remove_with_default
8.5e+02 ns pop
7.3e+02 ns remove_with_default