Store indexes after concatenating a numpy array-CodePudding

I have a list like this,

mylist = [
    np.array([48.5, 38.0, 40.0]),
    np.array([61.5, 52.5, 55.5, 46.5]),
    np.array([35.5, 36.5]),
]

I want to find the index of the array, and the location of the specific value in the array together with the values in mylist.

I am able to achieve the last column with np.concatenate(mylist) but don't know how to handle the rest efficiently.

expected = np.vstack(
    (
        np.array([0, 0, 0, 1, 1, 1, 1, 2, 2]),
        np.array([0, 1, 2, 0, 1, 2, 3, 0, 1]),
        np.array([48.5, 38.0, 40.0, 61.5, 52.5, 55.5, 46.5, 35.5, 36.5]),
    )
).T

It can be read as i.e. 38 is in the first array (index=0) and it is the second element of that array (index = 1).

CodePudding user response：

This does what you ask, if this is really what you want.

import numpy as np

mylist = [
    np.array([48.5, 38. , 40. ]),
    np.array([61.5, 52.5, 55.5, 46.5 ]),
    np.array([35.5, 36.5])]

a1 = []
a2 = []
for i,l in enumerate(mylist):
    a1.extend( [i] * len(l) )
    a2.extend( list(range(len(l))) )

final = np.array( [a1, a2, np.concatenate(mylist)] ).T
print(final)

Output:

[[ 0.   0.  48.5]
 [ 0.   1.  38. ]
 [ 0.   2.  40. ]
 [ 1.   0.  61.5]
 [ 1.   1.  52.5]
 [ 1.   2.  55.5]
 [ 1.   3.  46.5]
 [ 2.   0.  35.5]
 [ 2.   1.  36.5]]

CodePudding user response：

You can use use map with len to find out len of each sublist in mylist. Then use that in np.repeat to get "X" co-ordinate. Now, apply np.arange on each of the lengths to get "Y" co-ordinate and concatenate them using np.hstack. Now, just np.column_stack them together.

lens = list(map(len, mylist))
idx0 = np.repeat(np.arange(len(mylist)), lens) # [0, 0, 0, 1, 1, 1, 1, 2, 2]
idx1 = np.hstack([np.arange(v) for v in lens]) # [0, 1, 2, 0, 1, 2, 3, 0, 1]
vals = np.hstack(mylist) # [48.5, 38. , 40. , 61.5, 52.5, 55.5, 46.5, 35.5, 36.5]
out = np.column_stack([idx0, idx1, vals])

print(out)
[[ 0.   0.  48.5]
 [ 0.   1.  38. ]
 [ 0.   2.  40. ]
 [ 1.   0.  61.5]
 [ 1.   1.  52.5]
 [ 1.   2.  55.5]
 [ 1.   3.  46.5]
 [ 2.   0.  35.5]
 [ 2.   1.  36.5]]