cf2 cf2 - 3 months ago 14
Python Question

Combine list of numpy arrays and reshape

I'm hoping anybody could help me with the following.
I have 2 lists of arrays, which should be linked to each-other. Each list stands for a certain object.

arr1
and
arr2
are the attributes of that object.
For example:

import numpy as np

arr1 = [np.array([1, 2, 3]), np.array([1, 2]), np.array([2, 3])]
arr2 = [np.array([20, 50, 30]), np.array([50, 50]), np.array([75, 25])]


The arrays are linked to each other as in the
1
in
arr1
, first array belongs to the
20
in
arr2
first array. The result I'm looking for in this example would be a numpy array with size 3,4. The 'columns' stand for 0, 1, 2, 3 (the numbers in arr1) and the rows are filled with the corresponding values of arr2. When there are no corresponding values this cell should be 0.
Example:

array([[ 0, 20, 50, 30],
[ 0, 50, 50, 0],
[ 0, 0, 75, 25]])


How would I link these two list of arrays and reshape them in the desired format as shown in the above example?

Many thanks!

Answer

Here's an almost* vectorized approach -

lens = np.array([len(i) for i in arr1])

N = len(arr1)
row_idx = np.repeat(np.arange(N),lens)
col_idx = np.concatenate(arr1)

M = col_idx.max()+1
out = np.zeros((N,M),dtype=int)
out[row_idx,col_idx] = np.concatenate(arr2)

*: Almost because of the loop comprehension at the start, but that should be computationally negligible as it doesn't involve any computation there.