Home > Software engineering >  unzip list of tuples into numpy arrays
unzip list of tuples into numpy arrays

Time:11-04

Consider a list of tuples such as:

a = [(1,2), (3,4)]

I often find myself trying to unzip list like these into separate lists for each column value e.g.:

b,c = list(zip(*a))

In this case, b will be a list containing values 1 and 3.

I often find my self wanting b and c to be numpy arrays an not lists. In this case what I usually do is:

b,c = list(zip(*a))
b = np.array(b)
c = np.array(c)

The last two lines look cumbersome. Is there any way to unzip a list directly into two numpy arrays without casting them directly through numpy.array?

Thank you

CodePudding user response:

Your list of tuples can be converted into a 2-d numpy array by calling np.array. This can then be transposed and then unpacked along the first dimension using tuple assignment:

b, c = np.array(a).T

Here this gives:

>>> import numpy as np
>>> a = [(1,2), (3,4)]
>>> b, c = np.array(a).T   # or:  np.array(a).transpose()
>>> b
array([1, 3])
>>> c
array([2, 4])

Caveat: you will have a temporary array with the same number of elements as a, so it might be less memory-efficient than your original solution, particularly if you are unpacking into a larger number of 1d arrays.

  • Related