CODEWITHSUNDEEP
pythongrouping
Avra Sengupta
Avra Sengupta

Reputation: 11

Grouping data in Python

I have the following (space delimited) input:

2012-10-05 PETER 6  
2012-10-05 PETER 4  
2012-10-06 PETER 60
2012-10-06 TOM 10  
2012-10-08 SOMNATH 80

And I would like to achieve the following pipe-delimited output: (where the columns are [DATE AND NAME, NUM ENTRIES, SUM OF LAST COL])

2012-10-05 PETER|2|10  
2012-10-06 PETER|1|60  
2012-10-06 TOM|1|10  
2012-10-08 SOMNATH|1|80

This is my code so far:

s = open("output.txt","r")
fn=s.readlines()
d = {}
for line in fn:
 parts = line.split()
 if parts[0] in d:
   d[parts[0]][1] += int(parts[2])
   d[parts[0]][2] += 1
 else:
d[parts[0]] = [parts[1], int(parts[2]), 1]
for date in sorted(d):
   print "%s %s|%d|%d" % (date, d[date][0], d[date][2], d[date][1])

I am getting the output as:

2012-10-06 PETER|2|70

instead of 

2012-10-06 PETER|1|60

and TOM isn't showing in the list.

What do I need to do to correct my code?

Upvotes: 1

Views: 3563

Answers (2)

Janne Karila
Janne Karila

Reputation: 25207

d = collections.defaultdict(list)
with open('output.txt', 'r') as f:
    for line in f:
        date, name, val = line.split()
        d[date, name].append(int(val))

for (date, name), vals in sorted(d.items()):
    print '%s %s|%d|%d' % (date, name, len(vals), sum(vals))

Upvotes: 2

nneonneo
nneonneo

Reputation: 179717

<3 itertools

import itertools
with open('output.txt', 'r') as f:
    splitlines = (line.split() for line in f if line.strip())
    for (date, name), bits in itertools.groupby(splitlines, key=lambda bits: bits[:2]):
        total = 0
        count = 0
        for _, _, val in bits:
            total += int(val)
            count += 1
        print '%s %s|%d|%d' % (date, name, count, total)

If you don't want to use groupby (either it is unavailable, or your input data isn't guaranteed to be sorted), here's a conventional solution (which is effectively just a fixed version of your code):

d = {}
with open('output.txt', 'r') as f:
    for line in f:
        date, name, val = line.split()
        key = (date, name)
        if key not in d:
            d[key] = [0, 0]
        d[key][0] += int(val)
        d[key][1] += 1

for key in sorted(d):
    date, name = key
    total, count = d[key]
    print '%s %s|%d|%d' % (date, name, count, total)

Note that we use (date, name) as the key instead of just using date.

Upvotes: 0

Related Questions