what is the most pythonic way to split a 2d array to arrays of each row?

user_5120089 0
Viewed 8110

I have a function foo that returns an array with the shape (1000, 2) how can I split it to two arrays a(1000) and b(1000) I'm looking for something like this:

a;b = foo()

I'm looking for an answer that can easily generalize to the case in which the shape is (1000, 5) or so.

multidimensional-array arrays python
6 Answers

The zip(*...) idiom transposes a traditional more-dimensional Python list:

x = [[1,2], [3,4], [5,6]]

# get columns    
a, b = zip(*x)   # zip(*foo())
# a, b = map(list, zip(*x))  # if you prefer lists over tuples
a
# (1, 3, 5)

# get rows
a, b, c = x
a
# [1, 2]
user_23901820

Transpose and unpack?

a, b = foo().T

>>> a, b = np.arange(20).reshape(-1, 2).T
>>> a
array([ 0,  2,  4,  6,  8, 10, 12, 14, 16, 18])
>>> b
array([ 1,  3,  5,  7,  9, 11, 13, 15, 17, 19])
user_31255660

You can use numpy.hsplit.

x = np.arange(12).reshape((3, 4))
np.hsplit(x, x.shape[1])

This returns a list of subarrays. Note that in the case of a 2d input, the subarrays will be shape (n, 1). Unless you wrap a function around it to squeeze them to 1d:

def split_1d(arr_2d):
    """Split 2d NumPy array on its columns."""
    split = np.hsplit(arr_2d, arr_2d.shape[1])
    split = [np.squeeze(arr) for arr in split]
    return split

a, b, c, d = split_1d(x)

a
# array([0, 4, 8])

d
# array([ 3,  7, 11])
user_79545040

You could just use list comprehensions, e.g.

(a,b)=([i[0] for i in mylist],[i[1] for i in mylist])

To generalise you could use a comprehension within a comprehension:

(a,b,c,d,e)=([row[i] for row in mylist] for i in range(5))
user_75499070

You can do this simply by using zip function like:

def foo(mylist):
   return zip(*mylist)

Now call foo with as much dimension as you have in mylist, and it would do the requisite like:

mylist = [[1, 2], [3, 4], [5, 6]]
a, b = foo(mylist)
# a = (1, 3, 5)
# b = (2, 4, 6)
user_35355470

So this is a little nuts, but if you want to assign different letters to each sub-array in your array, and do so for any number of sub-arrays (up to 26 because alphabet), you could do:

import string
letters = list(string.ascii_lowercase)  # get all of the lower-case letters
arr_dict = {k: v for k, v in zip(letters, foo())}

or more simply (for the last line):

arr_dict = dict(zip(letters, foo()))

Then you can access each individual element as arr_dict['a'] or arr_dict['b']. This feels a little mad-scientist-ey to me, but I thought it was fun.

user_24645970
Related