Numpy在2D矩陣上的where（）

Question

我有一個像這樣的矩陣

t = np.array([[1,2,3,'foo'],
 [2,3,4,'bar'],
 [5,6,7,'hello'],
 [8,9,1,'bar']])

我想得到行包含字符串'bar'的索引

在1d數組中

rows = np.where(t == 'bar')

應該給我指數[0,3]然后廣播： -

results = t[rows]

應該給我正確的行

但我無法弄清楚如何讓它與2d數組一起工作。

Answer 1

對於一般情況，您的搜索字符串可以位於任何列中，您可以執行以下操作：

>>> rows, cols = np.where(t == 'bar')
>>> t[rows]
array([['2', '3', '4', 'bar'],
       ['8', '9', '1', 'bar']],
      dtype='|S11')

Answer 2

您必須將數組切片到要編制索引的列：

rows = np.where(t[:,3] == 'bar')
result = t1[rows]

返回：

 [[2,3,4,'bar'],
  [8,9,1,'bar']]