Мы можем использовать np.argpartition
, который делает indirect partition
и, таким образом, достигает некоторой эффективности там -
def nth_largest(prod): # works on prod from numpy.inner output
idx = np.argpartition(prod,-2,axis=1)[:,-2:]
I = np.arange(len(idx))
idx_s = prod[I[:,None],idx].argsort(1)
n_largest_indices = idx[I,idx_s[:,0]]
max_vals = prod[I,n_largest_indices]
return list(zip(I,n_largest_indices,max_vals))
Если вашей главной задачей является память, прибегните к циклу -
def innerprod_nth_largest_loopy(A, k): # works on input A
idxs = np.empty(len(A),dtype=np.uint64)
vals = np.empty(len(A),dtype=A.dtype)
for i,a in enumerate(A):
r = a.dot(A.T)
idx = np.argpartition(r,-k)[-k:]
idxs[i] = idx[r[idx].argsort()[0]]
vals[i] = r[idxs[i]]
return list(zip(range(len(A)),idxs,vals))
Заметьте, однако, что петлевая версия будет намного медленнее, просто хорошо со стороны памяти.