Adding & Removing values from pandas dataframe or CSV with reference dictionary

ihavealotofquestions00 0
Viewed 22

Fairly specific problem here - Below is the code i'm currently using that has a two inputs:

  • A reference dictionary (ref_dict in the code below) that has a single number as the keys, and the values being lists of 2-5 numbers (right now this is static in the file for building purposes, will build it out to pull from a file instead eventually)
  • A pandas df (right now being read from a csv for building purposes, but when implemented into production will be given an actual dataframe directly), that looks something like this:
week day country contsize code total
W1 Tu US M 1 8
W1 Tu CA M 2 4
W1 Tu US M 10 13
W1 Tu US M 20 2
W1 Tu CA M 20 16
W1 W US M 30 9

What i'm trying to do is:

  • Iterate through the rows of the df, look into the 'code' column and see if it matches any of the keys in the dictionary
  • If it matches any keys in the dictionary, I want to remove the 'total' and add it to the list of values related to the key in the dictionary ONLY if the other columns match. (so for example code 1 should give its total to to codes 10, 20, 30, code 2 should give it to 10, 20, etc. assuming all other columns are constant)
  • So in the table above example, the first row would take the 8, add it to row 3 and 4 (as code 10 and 20 are within the reference dictionary), but not 5 because the country is different, and not 6 because the day is different
  • the second row would take the 4 and would only add to row 5 (as code 20 is in the reference dictionary) but not any other as the country is different
  • Delete row 1 and 2

Expected output:

week day country contsize code total
W1 Tu US M 10 21
W1 Tu US M 20 10
W1 Tu CA M 20 20
W1 W US M 30 9

As of now with my below code, it is currently running through, adding totals and deleting the appropriate rows, except the totals are increasing by all of the values rather than the ones outlined in the reference dictionary.

Code - 10 Week 1 - Monday
Before 142
Expected 199
Actual After 221
Difference 22
Code 1 26
Code 2 26
Code 3 22
Code 4 5

With expected being 142 + 26 + 26 + 5 = 199 (Starting number + codes 1, 2, 4 as these 3 are the ones that should impact code 10) and difference being 221-199 = 22 (which is equal to #3 in the dict). I confirmed this across 2 separate weeks looking at all 5 possible days in the day column, its increasing by the exact amount.

Any tips?

import pandas as pd

def ref_dict():

    # Will update to pull from file instead
    ref_dict = {
        1: ['10', '20', '30'],
        2: ['10', '20'],
        3: ['20', '30'],
        4: ['10', '30']
    }
    return ref_dict

def unbundle(ref_dict):
    # Will update this to dataframe from IMS
    df = pd.read_csv('test.csv')

    df['remove'] = df.code.isin(ref_dict)

    for _, row in df[df.remove].iterrows():
        if row.code not in ref_dict:
            continue

        is_match = (df.week == row.week) & (df.day == row.day) & (df.country == row.country) & (df.contsize == row.contsize)
        df.loc[~df.remove & is_match, 'total'] += row.total

        df = df[~df.remove].drop(columns='remove')
        return df

def bundles():
    ref_dict = ref_dict() 
    new_totals = unbundle(ref_dict)
    new_totals.to_csv('final.csv')

bundles()
dataframe pandas python
1 Answers

Given:

# df
  week day country contsize  code  total
0   W1  Tu      US        M     1      8
1   W1  Tu      CA        M     2      4
2   W1  Tu      US        M    10     13
3   W1  Tu      US        M    20      2
4   W1  Tu      CA        M    20     16
5   W1   W      US        M    30      9

# ref_dict
{1: ['10', '20', '30'],
 2: ['10', '20'],
 3: ['20', '30'],
 4: ['10', '30']}

Doing:

# Create your mask:
remove = df.code.isin(ref_dict)

# Get the values from the dictionary:
df.loc[remove, 'code'] = df.loc[remove, 'code'].apply(ref_dict.get)

# Explode these values:
df = df.explode('code')

# Convert them to integers, since they're strings in your dict...
df.code = df.code.astype(int)

# Divide into two separate frames.
# Also, make them have an index of the values you need to match:
index = ['week', 'day', 'country', 'contsize', 'code']
remove_df = df.loc[remove].set_index(index)
df = df.loc[~remove].set_index(index)

# Sum the totals together, only those matching will be preserved,
# and reset the index back to how you had it:
df = df.assign(total=df.total.add(remove_df.total, fill_value=0)).reset_index()
print(df)

Output:

  week day country contsize  code  total
0   W1  Tu      US        M    10   21.0
1   W1  Tu      US        M    20   10.0
2   W1  Tu      CA        M    20   20.0
3   W1   W      US        M    30    9.0

Do not use iterrows or iter...anything. They're a trap and throw out the purpose of using pandas. If you absolutely need to iterate over something, use apply, but most of the time there will be a better method worth learning...

BeRT2me0
Related