Pandas 子集数据帧按索引条件

Question

我有一个如下所示的 DataFrame：

                   0
[0, 0, 0, 0, 0]  0.0    
[1, 0, 0, 0, 0]  1.0
...
[0, 1, 1, 1, 1]  0.0
[1, 1, 1, 1, 1]  1.0

我想将其子集为包含给定数量 1 的索引。

条件可能类似于

lambda x: np.count_nonzero(x) == n

n 个 1。

我知道可以使用 .loc 以某种方式按条件进行子集化，并且有很多答案可以解释在列上设置条件时如何进行。如果条件在索引本身上，我没有找到任何解释如何做到这一点的东西。

我试过了：

rounded.loc[np.count_nonzero(rounded.index) == n]

但没有运气。我不确定是否有一些非常明显的事情正在逃避我。

我可以创建一个新列，将索引作为值和基于它的子集，但我想知道是否有更优雅的解决方案。

Answer 1

在 pandas 中可以创建这个索引，但这并不容易。

对我来说，将index 转换为lists，然后再转换为array：

mask = np.count_nonzero(np.array(rounded.index.values.tolist()), axis=1) == n

示例：

df = pd.DataFrame({0:[0,1.0,0,1], 'a':[[0,0],[1,0],[1,1],[0,1]]})
rounded = df.set_index('a').rename_axis(None)
print (rounded)
          0
[0, 0]  0.0
[1, 0]  1.0
[1, 1]  0.0
[0, 1]  1.0

n = 1
mask = np.count_nonzero(np.array(rounded.index.values.tolist()), axis=1) == n
print (mask)
[False  True False  True]

print (rounded[mask])
          0
[1, 0]  1.0
[0, 1]  1.0