python - numpy 数组中列表的索引

Question

所以我正在编写一个代码，使用精确对角化来研究 Lieb-Liniger 模型。第一步是构建一个 numpy 数组，其中包含描述粒子占据的列表。该数组看起来像

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

对于 2 个粒子在 3 种模式下的情况。我的问题是，是否可以获取此数组中特定列表的索引，类似于使用 index 函数获取常规列表中的索引的方式。例如，对于列表列表 A，我能够使用 A.index(some_list_in_A) 获取该列表的索引，但我尝试使用 numpy.where(HS=[2,0,0]) 来获取 [2,0,0] 的索引(依此类推)，但无济于事。对于大量的粒子和模式，我正在寻找一种获取这些索引的有效方法，并且我认为使用 numpy 数组非常有效，但我刚刚遇到这个问题并没有找到解决方案。有什么建议吗？

Answer 1

你可以使用 np.where() 做:

pattern = [2,0,0]
index = np.where(np.all(a==pattern, axis=1))[0]