基于列值,将Pandas数据帧拆分为多个CSV文件

8

我有一个问题,非常类似于这个问题,但我需要更进一步,将分割后的数据帧保存为csv文件。

import pandas as pd
import numpy as np
import os

df = pd.DataFrame({ 'CITY' : np.random.choice(['PHOENIX','ATLANTA','CHICAGO', 'MIAMI', 'DENVER'], 1000),
                    'DAY': np.random.choice(['Monday','Tuesday','Wednesday', 'Thursday', 'Friday', 'Saturday', 'Sunday'], 1000),
                    'TIME_BIN': np.random.randint(1, 86400, size=1000),
                    'COUNT': np.random.randint(1, 700, size=1000)})

df['TIME_BIN'] = pd.to_datetime(df['TIME_BIN'], unit='s').dt.round('10min').dt.strftime('%H:%M:%S')
print(df)

OUTPUT:
         CITY  COUNT        DAY  TIME_BIN
0     ATLANTA    476   Thursday  12:20:00
1     PHOENIX     50   Saturday  15:40:00
2       MIAMI    250     Friday  08:20:00
3     CHICAGO    358     Monday  15:40:00
4     PHOENIX    217   Thursday  22:10:00
5       MIAMI     12   Thursday  21:40:00
6      DENVER     22     Friday  10:30:00
7     CHICAGO    645     Sunday  23:40:00
8       MIAMI    188     Sunday  08:40:00

我想为每个城市创建一个单独的数据框并将其保存为.csv文件。下面的代码可以工作,但如何在Pythonicly中执行而无需明确指定每个城市?真实数据集有大约20个城市,因此我不想重复这20次。我认为可以使用for循环在1-2行内完成下面的代码,但我不知道它会是什么样子。类似于 "for city in df['CITY']"。

df_phoenix = df[df['CITY'] == "PHOENIX"]
df_atlanta = df[df['CITY'] == "ATLANTA"]
df_chicago = df[df['CITY'] == "CHICAGO"]
df_phoenix.to_csv(os.getcwd() + "/data_phoenix.csv")
df_atlanta.to_csv(os.getcwd() + "/data_atlanta.csv")
df_chicago.to_csv(os.getcwd() + "/data_chicago.csv")
1个回答

17

我认为您需要使用自定义lambda函数或循环进行groupby:

f = lambda x: x.to_csv(os.getcwd() + "/data_{}.csv".format(x.name.lower()), index=False)
df.groupby('CITY').apply(f)

for i, x in df.groupby('CITY'):
     x.to_csv(os.getcwd() + "/data_{}.csv".format(i.lower()), index=False)

根据评论进行编辑,感谢@Anton vBR:

for i, x in df.groupby('CITY'):
    p = os.path.join(os.getcwd(), "data_{}.csv".format(i.lower()))
    x.to_csv(p, index=False)

1
谢谢,它能够正常工作并且是我一直在寻找的漂亮的一行代码。当我研究Pandas问题时,我经常在答案部分看到你的名字 :) - Calculus
1
不错,没想到可以在不使用循环的情况下实现。但是 os.getcwd() + 是必要的吗?如果我们不指定,默认情况下我们只需写入当前工作目录即可。 - sjw
@微积分一行代码可能很漂亮,但有时也不太易读。 - Anton vBR
2
@AntonvBR - 我添加了解决方案,你认为怎么样? - jezrael
1
优美的解决方案。 注意,os.path.join(os.getcwd(), )是多余的,可以删除。 - arielf

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接