Transform an array to 1s and 0s using another array-CodePudding

I have two arrays

arr1 = np.array([[4, 1, 3, 2, 5], [5, 2, 4, 1, 3]])
arr2 = np.array([[2], [1]])

I want to transform array 1 to a binary array using the elements of the array 2 in the following way

For row 1 of array 1, I want to use the row 1 of array 2 i.e. 2 - to make the top 2 values of array 1 as 1s and the rest as 0s
Similarly for row 2 of array 1, I want to use the row 2 of array 2 i.e. 1 - to make the top 1 value of array 1 as 1s and the rest as 0s

So arr1 would get transformed as follows

arr1_transformed = np.array([[1, 0, 0, 0, 1], [1, 0, 0, 0, 0]])

Here is what I tried.

arr1_sorted_indices = np.argosrt(-arr1)

This gave me the indices of the sorted array

array([[1, 3, 2, 0, 4],
       [3, 1, 4, 2, 0]])

Now I think I need to mask this array with the help of arr2 to get the desired output and I'm not sure how to do it.

CodePudding user response：

this should do the job in the mentioned case:

def trasform_arr(arr1,arr2):
  for i in range(0,len(arr1)):
    if i >= len(arr2):
      arr1[i] = [0 for x in arr1[i]]
    else:
      sorted_arr = sorted(arr1[i])[-arr2[i][0]:]
      arr1[i] = [1 if x in sorted_arr else 0 for x in arr1[i]]

arr1 = [[4, 1, 3, 2, 5], [5, 2, 4, 1, 3]]
arr2 = [[2], [1]]

trasform_arr(arr1,arr2)
print(arr1)

CodePudding user response：

You can try the following:

import numpy as np

arr1 = np.array([[4, 1, 3, 2, 5], [5, 2, 4, 1, 3]])
arr2 = np.array([[2],[1]])
 
s = np.argsort(np.argsort(arr1)[:, ::-1])
out = (np.arange(arr1.shape[1]) < arr2) * 1
out = out[np.arange(arr1.shape[0]).reshape(-1, 1), s]
print(out)

It gives:

[[1 0 0 0 1]
 [1 0 0 0 0]]