get the index of element in NumPy array
I have a Numpy integer array with a lot of duplicate elements.
For example:
a = np.random.randint(0,5,20)
a
Out[23]:
array([3, 1, 2, 4, 1, 2, 4, 3, 2, 3, 1, 4, 4, 1, 2, 4, 2, 4, 1, 1])

There are two cases:

if one element is less than 4, get all the indexes of this element
if one element is great than or equal to 4, select four of them randomly

I solved this with a loop.
ans = np.array([])
num = 4
for i in range(1,5):
    indexes = np.where(a == i)[0] # all indexes of elements equal to i
    index_i = np.random.choice(indexes, num, False) if len(indexes) &gt;=num else indexes
    ans = np.concatenate([ans, index_i])

np.sort(ans)

Out[57]:
array([ 0.,  1.,  2.,  5.,  6.,  7.,  8.,  9., 10., 11., 13., 14., 15.,
       17., 19.])

Can I solve this problem without a loop or more efficiently in Numpy or PyTorch?

You can do it quite easily, using Pandas.
First convert your array to a pandasonic Series:
s = pd.Series(a)

Then:

Group it by its value.
Apply to each group a function, which:

for groups of size 4 or smaller returns just this group,
for groups with more members, returns a random sample of 4 elements
from them.


Drop the 0-th level of the resulting index (added during grouping).
Sort by the (original) index, to bring back the original order (without
the dropped elements, for now we have original values with their
corresponding indices).
Return the index of the above result, as a Numpy array.

The code to do it is:
s.groupby(s).apply(lambda grp: grp if grp.size &lt;= 4 else grp.sample(4))\
    .reset_index(level=0, drop=True).sort_index().index.values

For a sample array containg:
array([2, 2, 1, 0, 1, 0, 2, 2, 2, 3, 0, 2, 1, 0, 0, 3, 3, 0, 2, 4])

the result is:
array([ 0,  2,  4,  5,  7,  9, 10, 11, 12, 14, 15, 16, 17, 18, 19])

To show that this result is correct, I repeated the source array,
with &quot;x&quot; marks below the elements at the returned indices.
array([2, 2, 1, 0, 1, 0, 2, 2, 2, 3, 0, 2, 1, 0, 0, 3, 3, 0, 2, 4])
       x     x     x  x     x     x  x  x  x     x  x  x  x  x  x

