Row-based dynamic column subset mean in Pandas

Viewed 66

I have a problem that can be visualized in the following way:

Our Cat The Home They Able
Alice 10 15 NaN 30 20 25
Bob 12 NaN 14 29 NaN 30
John NaN 9 NaN NaN NaN 20
Tyler 11 12 13 24 25 26

In general, there is numeric data assigned to each person (index) in each column, but there are empty spaces. I am wondering how to fill each NaN with the average for the same person for the columns with the same name length as the column with the missing value. In other words, how to combine fillna() and mean() with some custom logic for which columns are taken into consideration. The perfect result would be:

Our Cat The Home They Able
Alice 10 15 12.5 30 20 25
Bob 12 13 14 29 29.5 30
John 9 9 9 20 20 20
Tyler 11 12 13 24 25 26

With the numbers in bold being the averages for the same person for the same "column length".

Unfortunately, in my real life scenario there are hundreds of columns, so I cannot manually list the corresponding columns for each of them.

Thanks for all the help in advance.

2 Answers

you can try:

df = df.groupby(df.columns.map(len), axis =1).apply(lambda x: x.T.fillna(x.mean(1)).T)

OUTPUT:

        Our   Cat   The  Home  They  Able
Alice  10.0  15.0  12.5  30.0  20.0  25.0
Bob    12.0  13.0  14.0  29.0  29.5  30.0
John    9.0   9.0   9.0  20.0  20.0  20.0
Tyler  11.0  12.0  13.0  24.0  25.0  26.0

This seems to be working:

# create a df to hold the per-person means for each column name length
meandf = pd.DataFrame(index=df.index, columns=df.columns, dtype=float)

# find the unique column name lengths
lengths = set(len(i) for i in df.columns)

# iterate over the lengths and find take the mean for that chunk of the df
for l in lengths:
    subsetcols = df.columns[[len(col) == l for col in df.columns]]
    personmeans = df.loc[:, subsetcols].mean(axis=1)
    meandf.loc[personmeans.index, subsetcols] = personmeans

# write to the original df
df[df.isna()] = meandf

Result:

>>> df
        Our   Cat   The  Home  They  Able
Alice  10.0  15.0  12.5  30.0  20.0    25
Bob    12.0  13.0  14.0  29.0  29.5    30
John    9.0   9.0   9.0  20.0  20.0    20
Tyler  11.0  12.0  13.0  24.0  25.0    26

I used meandf as intermediate structure for holding the means (I couldn't figure out the indexing without it). In each cell, it contains the mean for each person for each column name length:

>>> meandf
        Our   Cat   The  Home  They  Able
Alice  12.5  12.5  12.5  25.0  25.0  25.0
Bob    13.0  13.0  13.0  29.5  29.5  29.5
John    9.0   9.0   9.0  20.0  20.0  20.0
Tyler  12.0  12.0  12.0  25.0  25.0  25.0
Related