Filter out data based on dynamic, ARRAYlike column values

user_10943886 0
Viewed 120

I have sample data which can be downloaded from here. What I want to do, to filter this data based on element_counts, cluster_choice_prob_k_fold, benchmark_probabilities columns for every individual.

This is the logic: element_counts column has an array-like structure [3, 2, 0, 0, 0] (in this case). It shows, how many elements I should select from top values of benchmark_probabilities from each top5 cluster_choice_prob_k_fold values. In this example, this array tells that:

  1. From the biggestcluster_choice_prob_k_fold group select 3 values based on their benchmark_probabilities
  2. From the second biggest cluster_choice_prob_k_fold group select 2 values based on their benchmark_probabilities
  3. From the third biggest cluster_choice_prob_k_fold group select 0 values based on their benchmark_probabilities
  4. Similar to 3
  5. Similar to 4.

So, my end result should look like this:

individual  cluster_choice_prob_k_fold  benchmark_probabilities element_counts
9710535 0.512776    0.163837    [3, 2, 0, 0, 0]
9710535 0.512776    0.0986      [3, 2, 0, 0, 0]
9710535 0.512776    0.085191    [3, 2, 0, 0, 0]
9710535 0.294674    0.050787    [3, 2, 0, 0, 0]
9710535 0.294674    0.037609    [3, 2, 0, 0, 0]

For other individual values, I have element_counts array which look like differently, but always with 5 elements and always element values adding up to 5. How can I do it with Pandas? Difficulty is, to write a function which can take into account different arrays and use it with groupby

pandas-groupby pandas python
2 Answers

Idea of solution is create helper column for number of sorted groups and filtering.

Solution working with same values of element_counts per each group by individual column, also there are always 5 values in lists.

First because sample data are in one column some preprocessing:

import ast

df = pd.read_excel('sample.xlsx')
c = df.columns[0].split(',')
df = df.iloc[:, 0].str.split(',', expand=True, n=3).applymap(lambda x: x.strip('"'))
df.columns = c
df['individual'] = df['individual'].astype(int)
c = ['cluster_choice_prob_k_fold','benchmark_probabilities']
df[c] = df[c].astype(float)
print (df.head())
   individual  cluster_choice_prob_k_fold  benchmark_probabilities  \
0     9710535                    0.512776                 0.163837   
1     9710535                    0.512776                 0.098600   
2     9710535                    0.512776                 0.085191   
3     9710535                    0.512776                 0.067577   
4     9710535                    0.512776                 0.065420   

    element_counts  
0  [3, 2, 0, 0, 0]  
1  [3, 2, 0, 0, 0]  
2  [3, 2, 0, 0, 0]  
3  [3, 2, 0, 0, 0]  
4  [3, 2, 0, 0, 0]

Then sorting per both columns by DataFrame.sort_values and convert string repr of lists to lists of integers by ast.literal_eval: 

df = df.sort_values(['cluster_choice_prob_k_fold','benchmark_probabilities'], ascending=False)
df['element_counts'] = df['element_counts'].apply(ast.literal_eval)

Then create counter of groups by GroupBy.transform and factorize:

df['g'] = df.groupby('individual', sort=False)['cluster_choice_prob_k_fold'].transform(lambda x: pd.factorize(x)[0]).astype(int)

Create new DataFrame with remove duplicates of individual by DataFrame.drop_duplicates, DataFrame.explode lists to new column and create column for groups with numpy.tile:

df1 = (df.drop_duplicates('individual')
         .explode('element_counts')[['individual','element_counts']]
         .assign(g = np.tile(np.arange(5), df['individual'].nunique())))
print (df1)
   individual element_counts  g
0     9710535              3  0
0     9710535              2  1
0     9710535              0  2
0     9710535              0  3
0     9710535              0  4

So is possible use DataFrame.merge with left join to original and remove non match groups with missing values by DataFrame.dropna:

df = (df.merge(df1, how='left', on=['individual','g'], suffixes=('','_'))
        .dropna(subset=['element_counts_']))
print (df.head(10))
   individual  cluster_choice_prob_k_fold  benchmark_probabilities  \
0     9710535                    0.512776                 0.163837   
1     9710535                    0.512776                 0.098600   
2     9710535                    0.512776                 0.085191   
3     9710535                    0.512776                 0.067577   
4     9710535                    0.512776                 0.065420   
5     9710535                    0.512776                 0.054764   
6     9710535                    0.512776                 0.048068   
7     9710535                    0.512776                 0.018973   
8     9710535                    0.294674                 0.050787   
9     9710535                    0.294674                 0.037609   

    element_counts  g element_counts_  
0  [3, 2, 0, 0, 0]  0               3  
1  [3, 2, 0, 0, 0]  0               3  
2  [3, 2, 0, 0, 0]  0               3  
3  [3, 2, 0, 0, 0]  0               3  
4  [3, 2, 0, 0, 0]  0               3  
5  [3, 2, 0, 0, 0]  0               3  
6  [3, 2, 0, 0, 0]  0               3  
7  [3, 2, 0, 0, 0]  0               3  
8  [3, 2, 0, 0, 0]  1               2  
9  [3, 2, 0, 0, 0]  1               2

And last use GroupBy.apply with DataFrame.head for filter rows per groups by added column, last removed both helper columns:

df = (df.groupby(['individual', 'cluster_choice_prob_k_fold'], 
                 group_keys=False, 
                 sort=False)
        .apply(lambda x: x.head(x['element_counts_'].iat[0]))
        .drop(['g','element_counts_'], axis=1))
print (df)
   individual  cluster_choice_prob_k_fold  benchmark_probabilities  \
0     9710535                    0.512776                 0.163837   
1     9710535                    0.512776                 0.098600   
2     9710535                    0.512776                 0.085191   
8     9710535                    0.294674                 0.050787   
9     9710535                    0.294674                 0.037609   

    element_counts  
0  [3, 2, 0, 0, 0]  
1  [3, 2, 0, 0, 0]  
2  [3, 2, 0, 0, 0]  
8  [3, 2, 0, 0, 0]  
9  [3, 2, 0, 0, 0]
user_29010020

Honestly, the solution with pandas.DataFrame.groupby does not come to my mind. However, I wrote a little bit brute force solution that works. If you want to give it a try, here it is (assume df being your sample data):

dfs = []

for element_count in df.element_counts.unique():
    element_count_list = element_count.strip('[').strip(']').split(', ')
    element_count_list_of_ints = [int(i) for i in element_count_list]
    top_probas = list(df.cluster_choice_prob_k_fold.unique())[:len(element_count_list_of_ints)]

    for n, prob in zip(element_count_list_of_ints, top_probas):
        df_part = df_element\
            .loc[(df.cluster_choice_prob_k_fold == prob) & (df.element_counts == element_count), :]\
            .sort_values(by='benchmark_probabilities', ascending=False)\
            .head(n)\
            .copy()
        dfs.append(df_part)

df_final = pd.concat(dfs, ignore_index=True)

The final data frame is following:

>>> print(df_final)
   individual  cluster_choice_prob_k_fold  benchmark_probabilities  \
0     9710535                    0.512776                 0.163837   
1     9710535                    0.512776                 0.098600   
2     9710535                    0.512776                 0.085191   
3     9710535                    0.294674                 0.050787   
4     9710535                    0.294674                 0.037609   

    element_counts  
0  [3, 2, 0, 0, 0]  
1  [3, 2, 0, 0, 0]  
2  [3, 2, 0, 0, 0]  
3  [3, 2, 0, 0, 0]  
4  [3, 2, 0, 0, 0]
user_71222720
Related