python 处理 .csv 文件, 按某列合并两个表

查了很多博客, 终于找到了比较好的方法, 不需要傻乎乎的遍历了.

问题描述:

aim.csv 文件中记录了一些 ID, source.csv 文件中记录了更多的 ID 和这些ID 对应的 key. 两个文件中的 ID 不完全相同
我希望能够(1)根据 source.csv 对 aim.csv 文件中的每一行 ID 添加相应的 key; (2)分别统计两个文件中 key 对应的 ID 数并保存在表 result.csv 中.

1. 读取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()
1
2
3
4
5
6
7
8
9
10
11
12
13

先去掉 aim 中的空值

这样, 需要用到的所有信息就先保存在 source 和 aim 两个 DataFrame 对象里面了. 此处用命令行演示, 直接创建 DataFrame 变量, 代替读取文件.

>>> 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
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

2. 根据 ID 查找 对应 key

>>> # 或 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
1
2
3
4
5
6
7

这一步中, 查找 source 中 ID 时直接用 .ID 和 ['ID'] 都行.
实际上是从 source 中挑选出了 ID 在 aim 中出现过的所有行, 因此顺序与 source 中 ID 排列顺序一致.

3. 按照 key 对相应的 ID 进行分组

>>> source = source.groupby('key')
>>> source
<pandas.core.groupby.generic.DataFrameGroupBy object at 0x000000000BD062B0>
>>> *source
  File "<stdin>", 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)
1
2
3
4
5
6
7
8
9
10
11
12
13

对 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]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

使用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
1
2
3
4
5
6
7
8

用 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
1
2
3
4
5
6
7

基本满足要求了, 但是为了和 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
1
2
3
4
5
6
7

对 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
1
2
3
4
5
6
7

3. 合并两个表

>>> 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
1
2
3
4
5
6
7

NaN 代表是空值.

4. 保存新表

result.to_csv('result.csv')
1

完成.

回头一看, 发现查找那一步就可以直接用 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
1
2
3
4
5
6
7

完.