I'm trying to create an 2d numpy array in shape n x k 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. Any help would be appreciated.

CodePudding user response：

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]])

CodePudding user response：

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)