python - 在列表理解中创建一个列表，然后在同一列表理解中从该列表创建另一个列表

Question

这个标题很拗口，所以可能更容易展示我试图通过代码实现的目标。

import pandas as pd

data = []
for i in range(0,6):
    row = []
    for j in range(0,6):
        if i == j:
            row.append(1)
        else:
            row.append(0)
    data.append(row)

colnames = [(1,2),(1,3),(1,4),(2,1),(3,1),(4,1)]
df = pd.DataFrame(data, columns = colnames)

以上不是我遇到问题的地方，但我想提供尽可能多的上下文。

我正在尝试迭代数据框并检索列名称的第一个元素，其中数据框的元素等于 1。我可以使用以下命令来执行此操作:

[col[0] for col in df if (df[col] == 1).any()]

这会生成我需要创建的第一个列表:

[1, 1, 1, 2, 3, 4]

我可以将该输出分配给变量并执行另一个列表理解以获得最终输出:

res = [col[0] for col in df if (df[col] == 1).any()]

[(res[i], res[i+1]) for i in range(len(res)-1)]

这会输出我的最终列表:

[(1, 1), (1, 1), (1, 2), (2, 3), (3, 4)]

是否可以在同一列表理解中执行这两项操作，同时仅接收最终列表作为输出？

这有点不优雅，但这就是非列表理解中的样子:

x = []
y = []
for i in range(0,1):
    for col in df:
        if df[col].any() == 1:
            x.append(col[0])
    for j in range(len(x)-1):
        y.append((x[j],x[j+1]))

感谢您花时间查看此内容!

Answer 1

您可以使用dot带有条件语句的产品，然后压缩并提取一个元组:

[(a,b) for a,b in zip(df.eq(1).dot(df.columns.str[0]),df.eq(1).dot(df.columns.str[0])[1:])]
#same with .iloc -> [(a,b) for a,b in zip(df.eq(1).dot(df.columns.str[0]),df.eq(1).dot(
#                                          df.columns.str[0]).iloc[1:])]

<小时/>

输出:

[(1, 1), (1, 1), (1, 2), (2, 3), (3, 4)]

完整工作代码:

data = []
for i in range(0,6):
    row = []
    for j in range(0,6):
        if i == j:
            row.append(1)
        else:
            row.append(0)
    data.append(row)

colnames = [(1,2),(1,3),(1,4),(2,1),(3,1),(4,1)]
df = pd.DataFrame(data, columns = colnames)
[(a,b) for a,b in zip(df.eq(1).dot(df.columns.str[0]),df.eq(1).dot(df.columns.str[0])[1:])]
#[(1, 1), (1, 1), (1, 2), (2, 3), (3, 4)]