python - python 中，如果前两个元素为 0 和 1，如何获取矩阵的行数？

Question

如果我有一个矩阵:

           A = [[0,1,2], [0,2,3], [1,5,6]] 

我想找到前两个元素为0和1的行，得到结果

           0

查找第一个元素为0的行，得到结果

           [0,1]

我该怎么办？最快的方法是什么？

Answer 1

我会这样做:

>>> A = [[0,1,2], [0,2,3], [1,5,6]]
>>> [i for i, row in enumerate(A) if row[:2] == [0, 1]]
[0]
>>> [i for i, row in enumerate(A) if row[0] == 0]
[0, 1]

如果您想同时创建两个结果，请使用常规循环，因为上面会迭代矩阵两次:

res_0, res_01 = list(), list()
for i, row in enumerate(A):
    if row[0] == 0:
        res_0.append(i)
        if row[1] == 1:
            res_01.append(i)