获取 NumPy 数组中元素的索引

Question

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:

1. if one element is less than 4, get all the indexes of this element
2. 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) >=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?

Answer 1

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 <= 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 "x" 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

Answer 2

Yes, you can do this using NumPy by:

a = np.random.randint(0,10,20)
print(a)

num = 4
if str(np.where(a<num)[0].shape) != '(0,)':             # Condition 1
    ans = np.where(a<num)[0]
    print(ans)
if str(np.where(a>=num)[0].shape) != '(0,)':            # Condition 2
    ans = np.random.choice(a[np.where(a>=num)[0]], 4)
    print(ans)

'''Output:
[9 9 8 1 0 7 7 4 6 2 8 2 1 2 9 5 5 1 4 1]
[ 3  4  9 11 12 13 17 19]
[4 9 8 7]
'''

I have done only for the cases you have mentioned. There can be many other cases such as if both conditions are true, or if there are less than 4 numbers in second case.