CODEWITHSUNDEEP
python
TobiasS
TobiasS

Reputation: 39

Sum an array with conditions in Python

I have a big array with data. I want to sum columns with one or two conditions. The data is already stored as classes in a dictionary.

The data is quite extensive, but the important part looks like this;

[["Gothenburg", "2018-01-05", "jan", 1.5, 2.3, 107],
 ["Gothenburg", "2018-01-15", "jan", 1.3, 3.3, 96],
 ["Gothenburg", "2018-01-25", "jan", 1.7, 3.2, 45],
 ["Gothenburg", "2018-03-05", "mar", 1.5, 2.1, 96],
 ["Gothenburg", "2018-03-05", "mar", 1.9, 2.8, 102],
 ["Malmo", "2018-01-02", "jan", 1.6, 2.3, 104],
 ["Malmo", "2018-01-10", "jan", 1.0, 2.9, 112],
 ["Malmo", "2018-03-05", "mar", 0.7, 4.3, 151],
 ["Malmo", "2018-03-25", "mar", 1.0, 3.3, 98],
 ["Hallsberg", "2018-01-25", "jan", 2.5, 2.3, 87],
 ["Hallsberg", "2018-02-14", "feb", 2.2, 2.3, 168],
 ["Hallsberg", "2018-03-06", "mar", 3.7, 2.3, 142],
 ["Hallsberg", "2018-04-29", "apr", 2.7, 2.3, 100]]

Explanation of columns: 0 = city, 1 = date, 2 = month, 3 = meanvalue1, 4 = meanvalue2, 5 = meanvalue3

The array is about 8000 rows in total with maybe 300 different cities.

What i want to achieve is to sum columns 3, 4, 5 after value in column 0, 1, 2.

For example sum of column 3 with key "Malmo" = 1.6 + 1.0 + 0.7 + 1.0 = 4.3 sum of column 3 with key "Malmo" and "jan" = 1.6 + 1.0 = 2.6

These conditional sums could either be stored in a dictionary (or a better solution), or they can be displayed att screen.

I guess there is a clever way to do this quite easy, but i haven't figured it out. I have tried to use for-loops and if cases, but it's messy. Hope to get some good advices here!

Upvotes: 0

Views: 197

Answers (2)

Parmandeep Chaddha
Parmandeep Chaddha

Reputation: 514

I like using the pandas library for dataframe-type objects. A solution for your problem:

import pandas as pd 
df  = pd.DataFrame([["Gothenburg", "2018-01-05", "jan", 1.5, 2.3, 107],
 ["Gothenburg", "2018-01-15", "jan", 1.3, 3.3, 96],
 ["Gothenburg", "2018-01-25", "jan", 1.7, 3.2, 45],
 ["Gothenburg", "2018-03-05", "mar", 1.5, 2.1, 96],
 ["Gothenburg", "2018-03-05", "mar", 1.9, 2.8, 102],
 ["Malmo", "2018-01-02", "jan", 1.6, 2.3, 104],
 ["Malmo", "2018-01-10", "jan", 1.0, 2.9, 112],
 ["Malmo", "2018-03-05", "mar", 0.7, 4.3, 151],
 ["Malmo", "2018-03-25", "mar", 1.0, 3.3, 98],
 ["Hallsberg", "2018-01-25", "jan", 2.5, 2.3, 87],
 ["Hallsberg", "2018-02-14", "feb", 2.2, 2.3, 168],
 ["Hallsberg", "2018-03-06", "mar", 3.7, 2.3, 142],
 ["Hallsberg", "2018-04-29", "apr", 2.7, 2.3, 100]])

df.columns = ['City', 'Date', 'Month', 'Mean1', 'Mean2', 'Mean3']

Choose what to group by:

group_by = ['City', 'Month'] #group_by = ['Month']

Create a group_by Dataframe with the sums of the columns:

City_Mon_Sum = df.groupby(group_by).agg({'Mean1': 'sum', 'Mean2': 'sum', 'Mean3': 'sum'}).reset_index()
City_Mon_Sum.rename(columns = {'Mean1': 'Group_Mean1', 'Mean2': 'Group_Mean2', 'Mean3': 'Group_Mean3'}, inplace = True )

Merge the two dataframes:

df = pd.merge(df, City_Mon_Sum, on = group_by)

Output: 

City    Date    Month   Mean1   Mean2   Mean3   Group_Mean1 Group_Mean2 Group_Mean3
0   Gothenburg  2018-01-05  jan 1.5 2.3 107           4.5   8.8          248
1   Gothenburg  2018-01-15  jan 1.3 3.3 96  4.5 8.8 248
2   Gothenburg  2018-01-25  jan 1.7 3.2 45  4.5 8.8 248
3   Gothenburg  2018-03-05  mar 1.5 2.1 96             3.4  4.9          198
4   Gothenburg  2018-03-05  mar 1.9 2.8 102 3.4 4.9 198
5   Malmo   2018-01-02  jan 1.6 2.3 104 2.6 5.2 216
6   Malmo   2018-01-10  jan 1.0 2.9 112 2.6 5.2 216
7   Malmo   2018-03-05  mar 0.7 4.3 151 1.7 7.6 249
8   Malmo   2018-03-25  mar 1.0 3.3 98  1.7 7.6 249
9   Hallsberg   2018-01-25  jan 2.5 2.3 87  2.5 2.3 87
10  Hallsberg   2018-02-14  feb 2.2 2.3 168 2.2 2.3 168
11  Hallsberg   2018-03-06  mar 3.7 2.3 142 3.7 2.3 142
12  Hallsberg   2018-04-29  apr 2.7 2.3 100 2.7 2.3 100

Upvotes: 1

Daniel Hepper
Daniel Hepper

Reputation: 29967

The trick is to use a tuple as key for the dictionary. Assuming your data is stored in a variable named big_array_with_data, here is a solution using collections.defaultdict:

from collections import defaultdict

monthly = [defaultdict(int) for i in range(3)]
totals =  [defaultdict(int) for i in range(3)]

for place, _, month, *means in big_array_with_data:
    for i, mean in enumerate(means):
        monthly[i][(place, month)] += mean
        totals[i][place] += mean

print(monthly[0][('Malmo', 'jan')])
print(totals[0]['Malmo'])

You could also do it without defaultdict like this:

monthly[i][(place, month)] = monthly[i].get((place, month), 0) + mean

That being said, if you are planning to do data crunching like this on a regular basis, working through a pandas tutorial is time well invested.

Upvotes: 0

Related Questions