Возьмите ненулевые элементы в макро-списке

Question

У меня возникла проблема с инструкцией np.nonzero() в python. Я хочу взять все индексы данного списка, которые не равны нулю. Итак, рассмотрим, что у меня есть следующий код:

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] 

, который соответствует элементам 1,4,5,6 от base, которые удовлетворяют этому условию. С другой стороны, я использовал команду:

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