在宏列表中取非零元素

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我想在base所有項目已在網站j,k非零分量，例如：

以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])