I think it is simple but it isn't working with me. I have 2 lists:
a = [1, 3, 6]
b = [['A', 'B', 'C', 'D', 'E', 'F', 'G'],
['H', 'I', 'J', 'K', 'L', 'M', 'N'],
['O', 'P', 'Q', 'R', 'S', 'T', 'U']]
and I need to iterate over b using a's elements.
The desired output is:
c = [['B', 'D', 'G'],
['I', 'K', 'N'],
['P', 'R', 'U']]
Arrays are welcome, Any suggestions?
You can use the itemgetter() method from the built in operator module:
from operator import itemgetter
a = itemgetter(1, 3, 6)
b = [['A', 'B', 'C', 'D', 'E', 'F', 'G'],
['H', 'I', 'J', 'K', 'L', 'M', 'N'],
['O', 'P', 'Q', 'R', 'S', 'T', 'U']]
c = [list(a(l)) for l in b]
print(c)
Output:
[['B', 'D', 'G'],
['I', 'K', 'N'],
['P', 'R', 'U']]
If a is supposed to index into each sublist you could use the following nested list comprehension
>>> [[sub[i] for i in a] for sub in b]
[['B', 'D', 'G'],
['I', 'K', 'N'],
['P', 'R', 'U']]
If a and b were numpy.array could also do
>>> b[:, a]
array([['B', 'D', 'G'],
['I', 'K', 'N'],
['P', 'R', 'U']], dtype='<U1')
(Assuming that the expected output just has a typo.)
The most efficient generic way is through operator.itemgetter().
from operator import itemgetter
a = [1, 3, 6]
b = [['A', 'B', 'C', 'D', 'E', 'F', 'G'],
['H', 'I', 'J', 'K', 'L', 'M', 'N'],
['O', 'P', 'Q', 'R', 'S', 'T', 'U']]
c = [list(itemgetter(*a)(x)) for x in b]
print(c)
# [['B', 'D', 'G'], ['I', 'K', 'N'], ['P', 'R', 'U']]
For your input sizes, the conversion to list brings it actually on par with a double list comprehension on my system, but it is otherwise faster. For larger a sizes, NumPy gets to actually be the fastest (although it requires the inner lists of b to have the same number of elements, which itemgetter does not care about).
Some timings as they come out from my test system:
%timeit [[x[i] for i in a] for x in b]
# 1000000 loops, best of 3: 1.14 µs per loop
%timeit [list(itemgetter(*a)(x)) for x in b]
# 1000000 loops, best of 3: 1.13 ns per loop
%timeit [itemgetter(*a)(x) for x in b]
# 1000000 loops, best of 3: 732 ns per loop
%timeit np.array(b)[:, tuple(a)]
# 100000 loops, best of 3: 6.84 µs per loop
b = b * 1000000
%timeit [[x[i] for i in a] for x in b]
# 1 loop, best of 3: 1.19 s per loop
%timeit [list(itemgetter(*a)(x)) for x in b]
# 1 loop, best of 3: 1.15 s per loop
%timeit [itemgetter(*a)(x) for x in b]
# 1 loop, best of 3: 800 ms per loop
%timeit np.array(b)[:, tuple(a)]
# 1 loop, best of 3: 2.31 s per loop
a = a * 100
b = b * 10000
%timeit [[x[i] for i in a] for x in b]
# 1 loop, best of 3: 386 ms per loop
%timeit [list(itemgetter(*a)(x)) for x in b]
# 10 loops, best of 3: 193 ms per loop
%timeit [itemgetter(*a)(x) for x in b]
# 10 loops, best of 3: 171 ms per loop
%timeit np.array(b)[:, tuple(a)]
# 10 loops, best of 3: 63.1 ms per loop
I like to use map to make operations on lists:
list(map(lambda x: [x[i] for i in a], b))
[['B', 'D', 'G'], ['I', 'K', 'N'], ['P', 'R', 'U']]
For its shortness, semantics and doing without additional modules I'd prefer a nested list comprehension:
>>> [[l[i] for i in a] for l in b]
[['B', 'D', 'G'], ['I', 'K', 'N'], ['P', 'R', 'U']]
The outer comprehension means 'create a list and add the inner comprehension's result for every list l in b to it.' The inner comprehension says 'create a list and add the i-th element of l for every i in a to it.'
I noticed that @norok2 tested how fast this option is, but didn't present it explicitly.
