查了很多博客, 终于找到了比较好的方法, 不需要傻乎乎的遍历了.
aim.csv
文件中记录了一些 ID, source.csv
文件中记录了更多的 ID 和这些ID 对应的 key. 两个文件中的 ID 不完全相同
我希望能够(1)根据 source.csv
对 aim.csv
文件中的每一行 ID 添加相应的 key; (2)分别统计两个文件中 key 对应的 ID 数并保存在表 result.csv
中.
使用 pandas
读取
import pandas
source_filename = "source.csv"
aim_filename = "aim.csv"
source_file = pandas.read_csv(source_filename)
aim_file = pandas.read_csv(aim_filename)
source = pandas.DataFrame({
'key':source_file['key'],
'ID':source_file['ID'],
})
aim = pandas.DataFrame({
'ID':aim_file['ID'],
})
aim = aim.dropna()
先去掉 aim
中的空值
这样, 需要用到的所有信息就先保存在 source
和 aim
两个 DataFrame
对象里面了. 此处用命令行演示, 直接创建 DataFrame
变量, 代替读取文件.
>>> import pandas
>>> aim = pandas.DataFrame({'ID':[0, 1, 2, 6]})
>>> source = pandas.DataFrame({'key':['a','b','b', 'c', 'd'],'ID':range(5)})
>>> aim
ID
0 0
1 1
2 2
3 6
>>> source
key ID
0 a 0
1 b 1
2 b 2
3 c 3
4 d 4
>>> # 或 aim = source[source['ID'].isin(aim['ID'])]
>>> aim = source[source.ID.isin(aim['ID'])]
>>> aim
key ID
0 a 0
1 b 1
2 b 2
这一步中, 查找 source
中 ID 时直接用 .ID
和 ['ID']
都行.
实际上是从 source
中挑选出了 ID 在 aim
中出现过的所有行, 因此顺序与 source
中 ID 排列顺序一致.
>>> source = source.groupby('key')
>>> source
<pandas.core.groupby.generic.DataFrameGroupBy object at 0x000000000BD062B0>
>>> *source
File "" , line 1
SyntaxError: can't use starred expression here
>>> print(*source)
('a', key ID
0 a 0) ('b', key ID
1 b 1
2 b 2) ('c', key ID
3 c 3) ('d', key ID
4 d 4)
对 source 按 key 分组, 得到一个 DataFrameGroupBy
对象. 这个对象操作起来有点麻烦
>>> source.agg(lambda x: len(list(x)))
ID
key
a 1
b 2
c 1
d 1
>>> source = source.agg(lambda x: list(x))
>>> source
ID
key
a [0]
b [1, 2]
c [3]
d [4]
使用aggregate方法, 重新获得了一个 DataFrame
对象, 获得了各个 key 对应的 ID 列表.
当然如果用 lambda x: len(list(x))
, 能直接获得对应的 ID 的数量, 但会丢失 ID 名的数据.
>>> source['source_IDs'] = source.apply(lambda x: len(x['ID']), axis = 1)
>>> source
ID source_IDs
key
a [0] 1
b [1, 2] 2
c [3] 1
d [4] 1
用 apply
方法添加新的一列, 值是对应的 ID 的数量, axis = 1 代表按列操作.
但是它的 index 即索引变成了 key, 因此需要重设一下.
>>> source = source.reset_index()
>>> source
key ID source_IDs
0 a [0] 1
1 b [1, 2] 2
2 c [3] 1
3 d [4] 1
基本满足要求了, 但是为了和 aim
表合并, 要给 ID 改个名.
>>> source = source.rename(columns = {'ID':'source_ID'})
>>> source
key source_ID source_IDs
0 a [0] 1
1 b [1, 2] 2
2 c [3] 1
3 d [4] 1
对 aim
也是一样的操作
>>> aim = aim.groupby('key').agg(lambda x: list(x))
>>> aim['aim_IDs'] = aim.apply(lambda x: len(x['ID']), axis = 1)
>>> aim = aim.reset_index().rename(columns = {'ID':'aim_ID'})
>>> aim
key aim_ID aim_IDs
0 a [0] 1
1 b [1, 2] 2
>>> result = source.merge(aim, how = 'left', on = ['key'])
>>> result
key source_ID source_IDs aim_ID aim_IDs
0 a [0] 1 [0] 1.0
1 b [1, 2] 2 [1, 2] 2.0
2 c [3] 1 NaN NaN
3 d [4] 1 NaN NaN
NaN 代表是空值.
result.to_csv('result.csv')
完成.
回头一看, 发现查找那一步就可以直接用 merge
,
>>> aim['key'] = aim.merge(source, how = 'left', on = 'ID')['key']
>>> aim
ID key
0 0 a
1 1 b
2 2 b
3 6 NaN
完.
对于学习pandas很有帮助
收藏夹里翻到了 玉树芝兰的 <<文科生数据科学上手指南>>, 从里面一路找到了 Python pandas Q&A video series byData School, 简直是手把手教学. 但是里面的 URL 好像都用不了, 毕竟加州都烧了…