CODEWITHSUNDEEP
pythonarrayslist
Flag
Flag

Reputation: 29

Python Iterating on list using another list

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?

Upvotes: 2

Views: 2701

Answers (5)

amadeusamadeus
amadeusamadeus

Reputation: 121

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.

Upvotes: 1

norok2
norok2

Reputation: 26956

(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

Upvotes: 1

Nicolas Gervais
Nicolas Gervais

Reputation: 36724

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']]

Upvotes: 1

Red
Red

Reputation: 27577

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']]

Upvotes: 2

Cory Kramer
Cory Kramer

Reputation: 118021

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')

Upvotes: 1

Related Questions