在宏列表

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