user2988577
user2988577

Reputation: 4207

Getting column from a multidimensional list

I have a quite involved nested list: each element is a tuple with two elements: one is an object, the other is an 3x2xn array. Here is a toy model.

toy=[('mol1',array([[[1,1,1],[2,2,2]],[[1,1,1],[2,2,2]]])),('mol2',array([[[1,1,1],[2,2,2]],[[1,1,1],[2,2,2]]]))]

How can I get a single column from that?

I am looking for

('mol1', 'mol2') 

and for the 2Darrays like:

array([[1,1,1],[1,1,1],[1,1,1],[1,1,1]])

I have a solution but I think it is pretty inefficient:

zip(*toy)[0]

it returns

('mol1', 'mol2')

then

zip(*toy)[1][0][:,0]

which returns

array([[1, 1, 1],
   [1, 1, 1]])

a for cycle like that

for i in range(len(toy)):
    zip(*toy)[1][i][:,0]

gives all the element of the column and I can build it with a vstack

Upvotes: 1

Views: 547

Answers (2)

piggybox
piggybox

Reputation: 1749

You can use the array in numpy to store your data or convert yours to that, then use the column slicing function built in. In general numpy slicing is very fast.

import numpy as np
np.asarray(toy)[::, 0] # first column

# output
array(['mol1', 'mol2'], 
      dtype='|S4')

Upvotes: 1

Roberto Bonvallet
Roberto Bonvallet

Reputation: 33339

This should be reasonably efficient:

>>> tuple(t[0] for t in toy)
('mol1', 'mol2')

For the 2D array, with the help of numpy's vstack function:

>>> from numpy import vstack
>>> vstack([t[1][:, 0] for t in toy])
array([[1, 1, 1],
       [1, 1, 1],
       [1, 1, 1],
       [1, 1, 1]])

Upvotes: 1

Related Questions