有关numpy的元素访问

Question

有一个二维的numpy初始数组：
s = array([[1,2],[3,2],[6,2],[4,1],[2,1],[9,2]])

x = np.where(s[:,1]==1)

x代表数组s中第二维等于1的元素的下标
所以x = (array([3, 4], dtype=int64),)

我的问题：已知下标,怎么快速创建一个np列表，每个元素是根据x来访问s的第一维的数，
比如上面返回的下标是[3,4]，我想要的结果是[s[3][0],s[4][0]]，这样。

如何通过简便的写法写出来?
我的尝试：

ans = [s[i][0] for i in x]，但是这样不对?为什么?
程序返回:[array([4, 1])]
而正确结果应该是:[array(4，2)]

Answer 1

import numpy as np

s = np.array([[1,2],[3,2],[6,2],[4,1],[2,1],[9,2]])
x = np.where(s[:, 1] == 1)

ans = [s[i, 0] for i in x[0]]
print(ans)

输出结果为: [4, 2]，即 [s[3][0], s[4][0]]