python输出二维数组中，每行N个最大值的索引

`import heapq
import numpy as np
import random
a = np.random.randint(50,size= (4,5))
a = np.array(a)
print(a)
lists = [[] for i in range(4)]
for i in range(len(a)):
# print(heapq.nlargest(3, range(len(a[i])), a[i].take))
lists[i].append(heapq.nlargest(3, range(len(a[i])), a[i].take))
print(lists)

`

输出结果如下：

[[14  5 29 32  2]
 [34 24 49 31 35]
 [39 37 43 40  5]
 [30  0 42 28  8]]

[[[3, 2, 0]], [[2, 4, 0]], [[2, 3, 0]], [[2, 0, 3]]]