Выбор определенных строк numpy.ndarray, содержащих numpy.ndarrays различного размера

Question

У меня есть numpy.ndarray, содержащий numpy.ndarray's разного размера, и я хочу извлечь несколько строк, в которых первое число каждой из этих строк находится в определенном списке.

Пример:

>>>a
array([array([1]), array([2, 3]), array([3, 1, 1]), array([2, 3, 3, 4])], dtype=object)

>>> idx
[2]

Мне нужно:

>>>a
array([array([2, 3]), array([2, 3, 3, 4])], dtype=object)

Так что я хочу только строки (или массивы), которые имеют 2 в качестве первого числа. Есть ли простой и приятный способ сделать это, как в этом посте здесь ?

Answer 1

Подобный массив объектов dtype больше похож на список списков, чем на 2d массив. a на самом деле имеет указатели на массивы элементов, просто список будет.

In [517]: a=array([array([1]), array([2, 3]), array([3, 1, 1]), array([2, 3, 3, 4])], dtype=object)
In [518]: a
Out[518]: 
array([array([1]), array([2, 3]), array([3, 1, 1]), array([2, 3, 3, 4])],
      dtype=object)

Простой способ выбора элементов a, соответствующих этому критерию, заключается в использовании понимания списка:

In [519]: [row for row in a if row[0]==2]
Out[519]: [array([2, 3]), array([2, 3, 3, 4])]

Почти все операции с массивом объекта типа, подобным этому, включают в себя итерацию списка, подобную этой.