Return only dataframe columns that satisfy where clause

user_6557593 0
Viewed 43

Starting with an arbitrary dataframe, I would like to return a dataframe with only those columns which have more than one distinct value.

I have:

X = df.nunique()

like:

    Id               5
    MSSubClass       3
    MSZoning         1
    LotFrontage      5
    LotArea          5
    Street           1
    Alley            0
    LotShape         2

Then I converted this from a series to a dataframe:

X = X.to_frame(name = 'dcount')

Then I used a where clause to only return values > 1:

X.where(X[['dcount']]>1)

which looks like:

                   dcount
    Id                5.0
    MSSubClass        3.0
    MSZoning          NaN
    LotFrontage       5.0
    LotArea           5.0
    Street            NaN
    Alley             NaN
    LotShape          2.0
    ...

But I now want only those column_names (in the index of X) which don't have dcount = 'NaN', so that I can ultimately go back to my original dataframe df and define it as:

df=df[[list_of_columns]]

How should this be done? I've tried a dozen ways and it's a PitA. I suspect there's a way to do it in 1 or 2 lines of code.

pandas python
1 Answers

You can use Boolean indexing and avoid converting your counts series to a dataframe:

counts = df.nunique()
df = df[counts[counts > 1].index]

The key is to note the index of your counts series are the column labels. So you can filter the series and then extract the required indices via pd.Series.index.

Here's a demonstration:

df = pd.DataFrame({'A': [1, 1, 1], 'B': [1, 2, 3],
                   'C': [4, 5, 5], 'D': [0, 0, 0]})

counts = df.nunique()
df = df[counts[counts > 1].index]

print(df)

   B  C
0  1  4
1  2  5
2  3  5
user_92095460
Related