oumilk
oumilk

Reputation: 1

Vectorizing Array of Values indexing another array

I am trying to vectorize the code below.

 for x in range (0, 500):
      S = rand.choice(np.unique(Y))   
      A = rand.choice(np.unique(X[Y==S])) 

Y and X are arrays where the values in the indices need to match.

Currently I modified S to be

 S = np.random.choice(np.unique(Y),size=500)

However, I'm not able to figure out how to index through X with an array of values S

An example for size = 5 could be

 Y = [0,0,2,3,2,4]
 X = [1,2,1,3,4,2]

 S = [0,2,0,3,2]
 X[Y==S] => ([1,2],[1,4],[1,2],[3],[1,4])  <=  Not sure how to get this
 A = [2,4,1,3,1]

Is there a simple way to do this?

Upvotes: 0

Views: 40

Answers (1)

andrew_reece
andrew_reece

Reputation: 21274

A little clumsy, but if you can use Pandas, combine isin() with a list comprehension:

import numpy as np
import pandas as pd

Y = pd.Series([0,0,2,3,2,4])
X = pd.Series([1,2,1,3,4,2])
S = pd.Series([0,2,0,3,2])

[np.random.choice(X[Y.isin([s])].values) for s in S]

You can get X[Y==S] with: [X[Y.isin([s])].values for s in S]

Upvotes: 1

Related Questions