Store indexes after concatenating a numpy array

Question

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).

It's not really clear what you are trying to do. Can you rewrite your question to show us your input and what you'd like the output to be. — Frank Yellin
– Frank Yellin, Commented Nov 30, 2021 at 7:01
@FrankYellin The expected is what I desire to have as an output. And the last sentence explains my goal. — datatech
– datatech, Commented Nov 30, 2021 at 7:03
Well, your first array will be [0]*len(mylist[0]) + [1]*len(mylist[1]) + [2]*len(mylist[2]), and the second is [range(len(mylist[0]))]+[range(len(mylist[1]))]+[range(len(mylist[2]))]. No need to go searching. — Tim Roberts
– Tim Roberts, Commented Nov 30, 2021 at 7:03

Tim Roberts · Accepted Answer · 2021-11-30 07:18:55Z

2

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]]

edited Nov 30, 2021 at 7:18

answered Nov 30, 2021 at 7:08

Tim Roberts

55.3k4 gold badges28 silver badges41 bronze badges

Sign up to request clarification or add additional context in comments.

4 Comments

Ch3steR Over a year ago

You've missed .T at the end to get expected value like the OP mentioned in the question.

Tim Roberts Over a year ago

Hmm, I missed the call to vstack. I will fix.

Ch3steR Over a year ago

No worries. Nice answer. You could write final = np.column_stack([a1, a2, np.concatenate(...)])

Tim Roberts Over a year ago

Since that's the essence of your answer, I'll leave it this way.

Ch3steR · Accepted Answer · 2021-11-30 07:08:31Z

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]]

Collectives™ on Stack Overflow

Store indexes after concatenating a numpy array

2 Answers 2

4 Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

4 Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related