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

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

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

Затем для каждой строки я случайным образом выбираю индекс столбца, который указывает значение больше 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

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


person gladys0313    schedule 11.05.2015    source источник
comment
Перебрать zip(*np.nonzero(a>2)).   -  person Ashwini Chaudhary    schedule 11.05.2015


Ответы (3)


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

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 или zip.

person Eli Korvigo    schedule 11.05.2015
comment
Конечно, с numpy намного проще, но и без него можно легко обойтись. - person Malik Brahimi; 11.05.2015
comment
Он уже загрузил numpy, мне кажется, что использовать встроенные методы еще проще и питоничнее. - person canyon289; 11.05.2015
comment
Я все еще получаю ValueError, распаковывая zip(*arr) в r,c - person Sean Summers; 11.05.2015
comment
Вы должны использовать np.vstack(np.where(arr>2)).T, а не zip. - person Ashwini Chaudhary; 11.05.2015
comment
@SeanSummers попробуйте точный код, который я разместил. В вашем коде может быть другая ошибка, которую вы не заметите. - person Eli Korvigo; 11.05.2015
comment
Я сделал, из свежей скорлупы. Еще слишком много значений для распаковки. Python 2.7.5, нумпи 1.9.1. - person Sean Summers; 11.05.2015
comment
@SeanSummers о, это была опечатка. Я использовал a в своем коде и заменил его на arr в ответе, но забыл заменить его в одном месте. Теперь это должно работать. Извини за это. - person Eli Korvigo; 11.05.2015
comment
@EliKorvigo все еще не работает (я отредактировал, чтобы включить свои версии). Ваш код работает для вас в новой оболочке? Какие версии вы используете? - person Sean Summers; 11.05.2015
comment
@SeanSummers должно было быть for r, c in zip(*indices), а не for r, c in zip(*a) - person Eli Korvigo; 11.05.2015
comment
Вы возродили мою веру в вырезание и склеивание, сэр! - person Sean Summers; 11.05.2015
comment
Программирование копирования-вставки @SeanSummers — это зло, как и переименование переменных вручную. Я очень извиняюсь за путаницу. - person Eli Korvigo; 11.05.2015
comment
@gladys0313 всегда пожалуйста. Кстати, у вас есть два ответа, которые работают, вы можете подумать о том, чтобы принять один из них. - person Eli Korvigo; 11.05.2015

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

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

person canyon289    schedule 11.05.2015

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

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]

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

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

person gladys0313    schedule 11.05.2015