Pandas subset dataframe by condition on indices

user_6334293 0
Viewed 503

I have a DataFrame that looks like this:

                   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

I would like to subset it to the indices that contain a given number of 1s.

The condition might look something like 

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

for n number of 1s.

I know that subsetting by condition can be done somehow with .loc, and there are many answers that explain how to do it when the condition is set on columns. I didn't find anything that explains how to do it if the condition is on the indices themselves.

I tried this:

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

But no luck. I am unsure whether something really obvious is escaping me.

I could just make a new column that has the indices as values and subset based on that, but I was wondering whether there is a more elegant solution.

pandas subset python
1 Answers

In pandas is possible create this index, but it is not easy.

For me working converting index to lists and then to array:

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

Sample:

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
user_29010020
Related