CODEWITHSUNDEEP
pythonmultidimensional-arraylist-comprehension
nodoze
nodoze

Reputation: 127

Select 2 columns from multidimensional array - list comprehension

I'm feeling extra dumb tonight. If I have a basic multi-dimensional list:

my_list = [['Bob', 43, 'Tall', 'Green'],
           ['Sally', 32, 'Short', 'Blue'],
           ['Tom', 54,'Medium','Orange']]

I can easily use a list comprehension to grab the first column:

new_list = [row[0] for row in my_list]

or I can grab 3 columns:

new_list = [row[0:3] for row in my_list]

but how would I grab columns 1 and 3 only?

new_list = [row[0,2] for row in my_list]

if not with a list comprehension, then how to accomplish with the least amount of code?

Upvotes: 0

Views: 2799

Answers (4)

Marcin
Marcin

Reputation: 238279

In addition to what @Simeon Visser wrote, you could also use itemgetter:

my_list = [['Bob', 43, 'Tall', 'Green'],
           ['Sally', 32, 'Short', 'Blue'],
           ['Tom', 54,'Medium','Orange']]

from operator import itemgetter

itg = itemgetter(0,2)

print([itg(row) for row in my_list])

Gives:

[('Bob', 'Tall'), ('Sally', 'Short'), ('Tom', 'Medium')]

Upvotes: 0

Mark Reed
Mark Reed

Reputation: 95267

While you can certainly use itemgetter - which you don't need to assign to a variable, nor do you need to fully qualify its name:

from operator import itemgetter
[itemgetter(0,2)(row) for row in my_list]

You could also just do what itemgetter does, which in this case would mean a nested list comprehension inside your list comprehension:

[[row[i] for i in [0,2]] for row in my_list]

Upvotes: 1

inspectorG4dget
inspectorG4dget

Reputation: 113975

import operator

In [7]: answer = [operator.itemgetter(0,2)(s) for s in my_list]

In [8]: answer
Out[8]: [('Bob', 'Tall'), ('Sally', 'Short'), ('Tom', 'Medium')]

Upvotes: 1

Simeon Visser
Simeon Visser

Reputation: 122376

One way:

new_list = [[row[0], row[2]] for row in my_list]

In this approach you're constructing a new list yourself for each row and you're putting element 0 and 2 in there as elements.

Upvotes: 2

Related Questions