Как получить индексы из списка/ndarray?

Question

У меня есть список, который выглядит как:

[[0,1,2], [1,2,3], [2,3,4], [3,4,5]] 

я могу сделать это в массив, как:

array([[0,1,2], 
     [1,2,3], 
     [2,3,4], 
     [3,4,5]]) 

Так все вместе у меня 4 строк, а каждая строка имеет 3 колонки. Теперь я хочу, чтобы найти индексы всех элементов, которые больше, чем 2, так и для всей матрицы, индексы должны быть:

((1,2),(2,1),(2,2),(3,1),(3,2),(3,3)) 

Тогда для каждой строки, я случайно выбрал индекс Col, который указывает значение больше 2. Теперь мой код, как:

a = np.array([[0,1,2],[1,2,3],[2,3,4],[3,4,5]] 
out = np.ones(4)*-1 
cur_row = 0 
col_list = [] 
for r,c in np.nonzero(a>2): 
    if r == cur_row: 
     col_list.append(c) 
    else: 
     cur_row = r 
     shuffled_list = shuffle(col_list) 
     out[r-1] = shuffled_list[0] 
     col_list = [] 
     col_list.append(c) 

Я надеюсь получить аут, который выглядит как:

array([-1, 2, 1, 2]) 

Однако, теперь, когда я запускаю мой код, он показывает

ValueError: too many values to unpack 

Кто-нибудь знает, как я исправить эту проблему? Или как мне сделать, чтобы достичь моей цели? Я просто хочу запустить код как можно быстрее, поэтому любые другие полезные идеи также более чем приветствуются.

Answer 1

Попробуйте это.

import numpy as np 

arr = np.array([[0,1,2], 
       [1,2,3], 
       [2,3,4], 
       [3,4,5]]) 
indices = np.where(arr>2) 

for r, c in zip(*indices): 
    print(r, c) 

Печать

1 2 
2 1 
2 2 
3 0 
3 1 
3 2 

Таким образом, он должен работать. Вы также можете использовать itertools.izip, в этом случае это был бы лучший выбор.

numpy Чистый раствор (благодаря @AshwiniChaudhary для предложения):

for r, c in np.vstack(np.where(arr>2)).T: 
    ... 

хотя я не уверен, что это будет быстрее, чем при использовании izip или почтовый индекс.

Answer 2

Вы можете просто сравнить массив с вашим значением и использовать где.

a = np.array([[0,1,2],[1,2,3],[2,3,4],[3,4,5]]) 
np.where(a>2) 

(массив ([1, 2, 2, 3, 3, 3], DTYPE = int64), массив ([2, 1, 2, 0, 1, 2], DTYPE = int64))

Чтобы получить ваши кортежи

list(zip(*np.where(a>2))) 

[(1, 2), (2, 1), (2, 2), (3, 0), (3, 1), (3, 2)]

Answer 3

Я сделал это, код должен быть:

a = np.array([[0,1,2],[1,2,3],[2,3,4],[3,4,5]]) 
out = np.ones(4)*-1 
cur_row = 0 
col_list = [] 
for r,c in zip(*(np.nonzero(a>2))): 
    if r == cur_row: 
     col_list.append(c) 
    else: 
     cur_row = r 
     shuffle(col_list) 
     if len(col_list) == 0: 
      out[r-1] = -1 
     else: 
      out[r-1] = col_list[0] 
     col_list = [] 
     col_list.append(c) 

shuffle(col_list) 
if len(col_list) == 0: 
    out[len(out)-1] = -1 
else: 
    out[len(out)-1] = col_list[0] 

Часть в конце концов, но за пределами forloop, чтобы убедиться, что последняя строка будет заботиться.

Это работает в моем случае.