CODEWITHSUNDEEP
python-3.xpandaspandas-groupby
Kumar AK
Kumar AK

Reputation: 1037

pandas remove records conditionally based on records count of groups

I have a dataframe like this 

import pandas as pd
import numpy as np

raw_data = {'Country':['UK','UK','UK','UK','UK','UK','UK','UK','UK','UK','UK','UK','UK','UK','UK','UK','UK','UK','UK','UK','UK'],
    'Product':['A','A','A','A','B','B','B','B','B','B','B','B','C','C','C','D','D','D','D','D','D'],
            'Week': [1,2,3,4,1,2,3,4,5,6,7,8,1,2,3,1,2,3,4,5,6], 
       'val': [5,4,3,1,5,6,7,8,9,10,11,12,5,5,5,5,6,7,8,9,10]
    }

df2 = pd.DataFrame(raw_data, columns = ['Country','Product','Week', 'val'])

print(df2)

enter image description here

and mapping dataframe 

mapping = pd.DataFrame({'Product':['A','C'],'Product1':['B','D']}, columns = ['Product','Product1'])

enter image description here

and i wanted to compare products as per mapping. product A data should match with product B data.. the logic is product A number of records is 4 so product B records also should be 4 and those 4 records should be from the week number before and after form last week number of product A and including the last week number. so before 1 week of week number 4 i.e. 3rd week and after 2 weeks of week number 4 i.e 5,6 and week 4 data.

similarly product C number of records is 3 so product D records also should be 3 and those records before and after last week number of product C. so product c last week number 3 so product D records will be week number 2,3,4.

wanted data frame will be like below i wanted to remove those yellow records

enter image description here

Upvotes: 1

Views: 55

Answers (2)

Valdi_Bo
Valdi_Bo

Reputation: 30971

Define the following function selecting rows from df, for products from the current row in mapping:

def selRows(row, df):
    rows_1 = df[df.Product == row.Product]
    nr_1 = rows_1.index.size
    lastWk_1 = rows_1.Week.iat[-1]
    rows_2 = df[df.Product.eq(row.Product1) & df.Week.ge(lastWk_1 - 1)].iloc[:nr_1]
    return pd.concat([rows_1, rows_2])

Then call it the following way:

result = pd.concat([ selRows(row, grp)
    for _, grp in df2.groupby(['Country'])
        for _, row in mapping.iterrows() ])

The list comprehension above creates a list on DataFrames - results of calls of selRows on:

  • each group of rows from df2, for consecutive countries (the outer loop),
  • each row from mapping (the inner loop).

Then concat concatenates all of them into a single DataFrame.

Upvotes: 1

jezrael
jezrael

Reputation: 862581

Solution first create mapped column by mapping DataFrame and create dictionaries for mapping for length and last (maximal) value by groups by Country and Product:

df2['mapp'] = df2['Product'].map(mapping.set_index('Product1')['Product'])
df1 = df2.groupby(['Country','Product'])['Week'].agg(['max','size'])

#subtracted 1 for last previous value
dprev = df1['max'].sub(1).to_dict()
dlen = df1['size'].to_dict()
print(dlen)
{('UK', 'A'): 4, ('UK', 'B'): 8, ('UK', 'C'): 3, ('UK', 'D'): 6}

Then Series.map values of dict and filter out less values, then filter by second dictionary by lengths with DataFrame.head:

df3 = (df2[df2[['Country','mapp']].apply(tuple, 1).map(dprev) <= df2['Week']]
            .groupby(['Country','mapp'])
            .apply(lambda x: x.head(dlen.get(x.name))))

print(df3)
                Country Product  Week  val mapp
Country mapp                                   
UK      A    6       UK       B     3    7    A
             7       UK       B     4    8    A
             8       UK       B     5    9    A
             9       UK       B     6   10    A
        C    16      UK       D     2    6    C
             17      UK       D     3    7    C
             18      UK       D     4    8    C

Then filter original rows unmatched mapping['Product1'], add new df3 and sorting:

df = (df2[~df2['Product'].isin(mapping['Product1'])]
           .append(df3, ignore_index=True)
           .sort_values(['Country','Product'])
           .drop('mapp', axis=1))
print(df)
   Country Product  Week  val
0       UK       A     1    5
1       UK       A     2    4
2       UK       A     3    3
3       UK       A     4    1
7       UK       B     3    7
8       UK       B     4    8
9       UK       B     5    9
10      UK       B     6   10
4       UK       C     1    5
5       UK       C     2    5
6       UK       C     3    5
11      UK       D     2    6
12      UK       D     3    7
13      UK       D     4    8

Upvotes: 0

Related Questions