如何从元素是整数的ndarray创建一个整数numpy 2darray索引？

Question

I'm trying to create an 2d numpy array in shape nxk where n is the dimension of the ndarray given and k is the amount of elements from the ndarray that are integers. Each row in the returned array should contain the indexes at which the condition holds at the relevant dimension. For example, the ndarray is:

array([[ 0.        , -0.36650892, -0.51839849,  4.55566517,  4.        ],
       [ 5.21031078,  6.29935488,  8.29787346,  7.03293348,  8.74619707],
       [ 9.36992033, 11.        , 11.88485714, 12.98729128, 13.98447014],
       [14.        , 16.71828376, 16.15909201, 17.86503506, 19.12607872]])

Again, the condition is if the element is an integer so the returned array should be:

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

Note that for the 0th row we want the 0th and 4th elements so we get [0,0,....],[0,4,...] and so on. I thought about creating a new array at the same shape as arr with True at the integer element positions and False elsewhere. Not sure where to proceed with this though.

Answer 1

Assuming a the input array, you can compare to the rounded values to identify the integers, use numpy.where to get their indices and np.vstack to form the final array:

np.vstack(np.where(a==a.round()))

output:

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

Answer 2

You can do something like this:

import numpy as np
a = np.array([[ 0.        , -0.36650892, -0.51839849,  4.55566517,  4.        ],
       [ 5.21031078,  6.29935488,  8.29787346,  7.03293348,  8.74619707],
       [ 9.36992033, 11.        , 11.88485714, 12.98729128, 13.98447014],
       [14.        , 16.71828376, 16.15909201, 17.86503506, 19.12607872]])

# check where the integer of a value is equal the value
mask = np.int_(a) == a

# get indexes where the mask is true
where = np.where(mask)

# make the output into an array with the shape you wanted
output = np.stack([where[0], where[1]])

print(output)

Output:

array([[0, 0, 2, 3],
       [0, 4, 1, 0]], dtype=int64)