在宏列表中获取非零元素

Question

我在python中的指令np.nonzero()有问题。 我想获取给定列表的所有非零索引。 因此，请考虑以下代码：

import numpy as np
from scipy.special import binom
M=4
N=3
def generate(N,nb):
        states = np.zeros((int(binom(nb+N-1, nb)), N), dtype=int)
        states[0, 0]=nb
        ni = 0  # init
        for i in xrange(1, states.shape[0]):
            states[i,:N-1] = states[i-1, :N-1]
            states[i,ni] -= 1
            states[i,ni+1] += 1+states[i-1, N-1]
            if ni >= N-2:
                if np.any(states[i, :N-1]):
                    ni = np.nonzero(states[i, :N-1])[0][-1]
            else:
                ni += 1
        return states
base = generate(M,N)

base的结果由下式给出：

    base = [[3 0 0 0]
 [2 1 0 0]
 [2 0 1 0]
 [2 0 0 1]
 [1 2 0 0]
 [1 1 1 0]
 [1 1 0 1]
 [1 0 2 0]
 [1 0 1 1]
 [1 0 0 2]
 [0 3 0 0]
 [0 2 1 0]
 [0 2 0 1]
 [0 1 2 0]
 [0 1 1 1]
 [0 1 0 2]
 [0 0 3 0]
 [0 0 2 1]
 [0 0 1 2]
 [0 0 0 3]]

关键是，对于给定的索引j,k我想将所有在j,k站点中具有非零成分的项目都作为base ，例如：

取j=0,k=1我必须获得：

result = [1 4 5 6]

对应于满足此条件的base元素1,4,5,6。 另一方面，我使用了以下命令：

np.nonzero((base[:, j]) & (base[:, k]))[0]

但是它不能正常工作，为什么？

Answer 1

首先，列表索引base[:, j]的语法错误，请使用： [:][j]代替

也：

np.nonzero((base[:, j]) & (base[:, k]))[0]

无法使用，因为&符号不适用于此处。

您可以这样使用numpy：

b = np.array(base);
j=0;k=1;
np.nonzero(b.T[j]* b.T[k])[0]

这将给：

array([1, 4, 5, 6])