Python 使用df.itertuples（）中的元组，如何在条件下检索每个元组元素的列值？

我有一个

pandas.DataFrame

，例如：

   1  2  3
1  1  0  0
2  0  1  0
3  0  0  1

从包含以下关系的集合中创建的：

{(1,1),(2,2),(3,3)}

我正在尝试为此创建等价类。大概是这样的：

[1] = {1}
[2] = {2}
[3] = {3}

到目前为止，我已经做了以下工作：

testGenerator = generatorTest(matrix)
indexCount = 1
while True:
    classRelation, loopCount = [], 1
    iterable = next(testGenerator)
    for i in iterable[1:]:
        if i == 1:
            classRelation.append(loopCount)
        loopCount += 1
    print ("[",indexCount,"] = ",set(classRelation))
    indexCount += 1

正如你所看到的，这是非常混乱的。但我确实得到了或多或少的期望输出：

[ 1 ] =  {1}
[ 2 ] =  {2}
[ 3 ] =  {3}

---

如何以更整洁、更具Python风格的方式完成相同的输出？

考虑数据帧

df

 df = pd.DataFrame(np.eye(3, dtype=int), [1, 2, 3], [1, 2, 3])

numpy.where

i, j = np.where(df.values == 1)
list(zip(df.index[i], df.columns[j]))

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

s = df.stack()
s.compress(s.astype(bool)).index.tolist()

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

堆栈

和

压缩

i, j = np.where(df.values == 1)
list(zip(df.index[i], df.columns[j]))

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

s = df.stack()
s.compress(s.astype(bool)).index.tolist()

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

在这种情况下，您可以像这样使用：

代码：

df.idxmax(axis=1)

df = pd.DataFrame([[1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 1, 0]],
                  columns=[1, 2, 3], index=[1, 2, 3, 4])
print(df.idxmax(axis=1))

1    1
2    2
3    3
4    2
dtype: int64

测试代码：

df.idxmax(axis=1)

df = pd.DataFrame([[1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 1, 0]],
                  columns=[1, 2, 3], index=[1, 2, 3, 4])
print(df.idxmax(axis=1))

1    1
2    2
3    3
4    2
dtype: int64

结果：

df.idxmax(axis=1)

df = pd.DataFrame([[1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 1, 0]],
                  columns=[1, 2, 3], index=[1, 2, 3, 4])
print(df.idxmax(axis=1))

1    1
2    2
3    3
4    2
dtype: int64

---

也许有一种误解。假设其中一行有多个

1

。例如，结果将只返回第一个

1

的位置，而不是所有位置。如果另一个答案不符合你的需求，你可以考虑扩大你的样本。