Find all rows indices in which elements of 2D NumPy array occur-CodePudding

Lets say we have 2D array like this one:

input_array = np.array([[0, 4, 6],
                        [5, 4, 1],
                        [2, 1, 0],
                        [4, 1, 0],
                        [1, 5, 3]])

How to get 2D array in which we have all the indices of rows in which element occurs set to one. Result should look like this:

output_array = [[1, 0, 1, 1, 0], # zero is in 0th, 2nd and 3rd row
                [0, 1, 1, 1, 1],
                [0, 0, 1, 0, 0],
                [0, 0, 0, 0, 1],
                [1, 1, 0, 1, 0],
                [0, 1, 0, 0, 1],
                [1, 0, 0, 0, 0]]

Simplification: We can assume that input_array has all the integer values in some range. None of them are missing. In this case: 0-6

CodePudding user response：

I figured it out:

rows = input_array.max()   1
cols = input_array.shape[0]

output_array = np.zeros((rows, cols))

indices = input_array * cols   np.arange(cols)[:, np.newaxis]
output_array = output_array.ravel()
output_array[indices.ravel()] = 1
output_array = output_array.reshape((rows, cols))

CodePudding user response：

Here's maybe a slightly more elegant solution using numpy's any() and a list comprehension. This also works if there are missing values in the range from 0 to 6.

>>> np.array([np.any(input_array == i, axis=1) for i in range(6)], dtype=int)

array([[1, 0, 1, 1, 0],
       [0, 1, 1, 1, 1],
       [0, 0, 1, 0, 0],
       [0, 0, 0, 0, 1],
       [1, 1, 0, 1, 0],
       [0, 1, 0, 0, 1]])