如何使用itertools.groupby将CSV数据转换为字典

时间:2019-11-23 05:45:45

标签: python dictionary itertools

我有一个文本文件job.txt,位于下面

job,salary
Developer,29000
Developer,28000
Tester,27000
Tester,26000

我的代码是

with open(r'C:\Users\job.txt') as f:
    file_content = f.readlines()
data = {}
for i, line in enumerate(file_content):
    if i == 0:
        continue
    job, salary = line.split(",")
    job = job.strip()
    salary = int(salary.strip())
    if not job in data:
        data[job] = []
    data[job].append(salary)
print("data =", data)

我的预期结果低于

data = {'Developer': [29000, 28000], 'Tester': [27000, 26000]}

如何将我的代码转换为使用itertools.groupby

3 个答案:

答案 0 :(得分:4)

如果您的数据已经按类别划分,则只能依靠groupby

from itertools import groupby

with open("job.txt") as f:
    rows = [x.split(",") for x in f.readlines()[1:]]

data = {
    k.strip(): [int(y[1]) for y in v]
    for k, v in groupby(rows, key=lambda x: x[0])
}

考虑到这一点,我认为defaultdict在这里更合适。订单是自动处理的,它不太聪明。此外,无需将文件存入内存或对其进行排序(如果未排序)。如果您不喜欢dict(data)子类,请在最后使用defaultdict

from collections import defaultdict

data = defaultdict(list)

with open("job.txt") as f:
    for i, line in enumerate(f):
        if i:
            job, salary = [x.strip() for x in line.split(",")]
            data[job].append(int(salary))

如已接受的答案中所述,如果您的实际数据比示例中的数据要复杂得多,请选择CSV module。 CSV可能很难解析,没有理由重新发明轮子。

答案 1 :(得分:2)

这是将生成所需字典的代码。

from itertools import groupby

data = [
    ["Developer",29000],
    ["Developer",28000],
    ["Tester",27000],
    ["Tester",26000]
]

def keyfunc(e):
    return e[0]

unique_keys = {}
data = sorted(data, key=keyfunc)

for k, g in groupby(data, keyfunc):
    unique_keys[k] = [i[1] for i in g]


>>> print(unique_keys)
{'Developer': [29000, 28000], 'Tester': [27000, 26000]}

P.S:我建议您使用csv模块来读取文件,而不要自己做。

答案 2 :(得分:2)

如果可以选择熊猫,请尝试以下操作:

from collections import defaultdict
import pandas as pd

d = pd.read_csv('job.txt').to_numpy().tolist() 
res = defaultdict(list)
for v, k in d: res[v].append(k)
d = dict(res)

d
# {'Developer': [29000, 28000], 'Tester': [27000, 26000]}