A=np.array([        [7,8],[7,9],[3,4],[5,4],[3,4],[5,6]    ])

indicesB=np.array([ [1]  ,[1]  ,[1]  ,[2]  ,[1]  ,[2]  ])

how can i get all the elements in A if the same position elements in indices B= 1?

for example,

if i want indicesB= 2,then i get[5,4],[5,6] if i want indicesB= 1,then i get[7,8],[7,9],[3,4],[3,4]

What I want is something like this

Y=np.array([[7,8],[3,4],[3,4],[3,4],[3,4],[3,4]])

X=np.array([[1],[1],[1],[1],[1],[2]])

for x in range(1,3):
 for i in range(6):
  if X[i]==x:
    print('the indice is ', x,Y[i])

how cccan i make it simple using numpy?

CodePudding user response：

If I understand right this code might helps you:

new_list = []
for i,j in zip(A, indicesB):
    el = i[:j[0]].tolist()
    new_list.append(el)

If you need an array instead of list you should use i[:j[0]] without .tolist() and after loop change type new_list to array like that:

new_list = np.array(new_list)

CodePudding user response：

Can you use dict ? that way you can call perticular key dict[1] and you will receive np.array.

import numpy as np
dic = {s:[] for s in set(np.concatenate([x.ravel() for x in X]))}
[dic[j.tolist()[0]].append(i.tolist()) for i,j in zip(A, B)]
    
np.array(dic[2]) #or np.array(dic[int(input())]) 

Output:

array([[5, 4],
       [5, 6]])