itertools.gropby使用

2016-11-18  本文已影响11人  enomine

问题背景

本周主要写了各种类型的财务报表,利用python的csv模块进行输出。

在编写财务报表的过程中,出现了这样一种需求。

姓名 月份 电费 水费 燃气费
张三 1 100.00
张三 1 30.00
张三 1 40.00
张三 1 10.00
张三 2 70.00
李四 3 66.00
李四 3 20.00

类似于这样的报表,要将姓名和月份一致的信息,合并到一条中去。结果像这样:

姓名 月份 电费 水费 燃气费
张三 1 100.00 30.00 40.00
张三 2 70.00
李四 3 66.00 20.00

我们使用csv模块 的 DictWriter

表1对应的数据结构为:

[{
    'name': '张三',
    'month': 1,
    'electricity_fee': 100.00,
}, {
    'name': '张三',
    'month': 1,
    'water_fee': 30.00,
}, {
    'name': '张三',
    'month': 1,
    'gas_fee': 40.00,
}, {
    'name': '张三',
    'month': 1,
    'water_fee': 10.00,
}, {
    'name': '张三',
    'month': 2,
    'water_fee': 70.00,
}, {
    'name': '李四',
    'month': 3,
    'electricity_fee': 66.00,
}, {
    'name': '李四',
    'month': 3,
    'gas_fee': 20.00,
}]

表2:

[{
    'name': '张三',
    'month': 1,
    'electricity_fee': 100.00,
    'water_fee': 40.00,
    'gas_fee': 40.00,
}, {
    'name': '张三',
    'month': 2,
    'water_fee': 70.00,
}, {
    'name': '李四',
    'month': 3,
    'electricity_fee': 66.00,
    'gas_fee': 20.00,
}]

如何用简介优雅的代码将表1的数据结构转化为表2?

from itertools import groupby

def get_key(data):
    return data['name'] + str(data['month'])

fees = ('electricity_fee', 'water_fee', 'gas_fee')

rows_out = []
for key, group in groupby(rows, key=get_key):
    row = group[0]
    for fee in fees:
        row[fee] = sum(item[fee] for item in group if fee in item)
        rows_out.append(row)
上一篇下一篇

猜你喜欢

热点阅读