Поведение np.where () для 2d массивов

Question

У меня есть массивный двумерный массив, и каждый элемент массива содержит словарь.

[[{'foo': 1} {'bar': 2, 'xyz': 7} {} {}]
 [{} {'xyz': 7} {'bar': 2} {'foo': 1}]
 [{} {'xyz': 7} {'foo': 1} {'bar': 2}]]

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

Мое текущее решение выглядит так:

indices = []
for row in arr:
  for i in range(len(row)):
    if 'foo' in row[i].keys():
      indices.append(i)

и возвращает

[0, 3, 2]

Но мне интересно, есть ли лучший способ написать это также с использованием np.where ()

Я обнаружил, что это почти то, что я ищу

np.where([[['foo' in ele.keys()] for ele in row] for row in arr])

, но он возвращает 3 массива (средний правильный)

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

Я не понимаю первый и третий массив, почему они созданы?

Answer 1

In [226]: [[['foo' in ele.keys()] for ele in row] for row in arr]                                    
Out[226]: 
[[[True], [False], [False], [False]],
 [[False], [False], [False], [True]],
 [[False], [False], [True], [False]]]

Сделать это массивом, дающим форму (3,4,1). where возвращает набор массивов, по одному для каждого измерения.

Использование группировки () вместо [] дает 2d массив:

In [227]: [[('foo' in ele.keys()) for ele in row] for row in arr]                                    
Out[227]: 
[[True, False, False, False],
 [False, False, False, True],
 [False, False, True, False]]

frompyfunc - это еще один способ применения функции к каждому элементу массива. Он имеет тенденцию быть немного быстрее, чем явные циклы (до 2х), и особенно хорош при работе с массивами dtype объекта:

In [228]: np.frompyfunc(lambda d: 'foo' in d.keys(),1,1)(arr)                                        
Out[228]: 
array([[True, False, False, False],
       [False, False, False, True],
       [False, False, True, False]], dtype=object)
In [229]: np.where(_)                                                                                
Out[229]: (array([0, 1, 2]), array([0, 3, 2]))

Answer 2

Вы также можете использовать это:

x = [[{'foo': 1},{'bar': 2, 'xyz': 7},{},{}],
     [{} ,{'xyz': 7} ,{'bar': 2} ,{'foo': 1}],
     [{}, {'xyz': 7}, {'foo': 1}, {'bar': 2}]]

res = [list(map(lambda elem: 'foo' in elem.keys(),item)) for item in x]
print(res)
print(numpy.where(res))

выход:

[[True, False, False, False],
 [False, False, False, True],
 [False, False, True, False]]

(array([0, 1, 2], dtype=int64), array([0, 3, 2], dtype=int64))