Вот векторизованный с masking
, а также с использованием broadcasting
-
def select_gt_indices(a, idx):
r = np.arange(a.shape[1])
select_mask = idx[:,None] <= r
put_mask = (a.shape[1]-idx-1)[:,None] >= r
# or put_mask = np.sort(select_mask,axis=1)[:,::-1]
out = np.zeros_like(a)
out[put_mask] = a[select_mask]
return out
Пробный прогон -
In [92]: np.random.seed(0)
...: a = np.random.randint(0,999,(4,5))
...: idx = np.array([2,4,3,0])
In [93]: a
Out[93]:
array([[684, 559, 629, 192, 835],
[763, 707, 359, 9, 723],
[277, 754, 804, 599, 70],
[472, 600, 396, 314, 705]])
In [94]: idx
Out[94]: array([2, 4, 3, 0])
In [95]: select_gt_indices(a, idx)
Out[95]:
array([[629, 192, 835, 0, 0],
[723, 0, 0, 0, 0],
[599, 70, 0, 0, 0],
[472, 600, 396, 314, 705]])