如何用Python对CSV文件的多个列进行排序？

Question

我在处理一个大数据集时遇到了一些问题，想把这些数据整理得更有用。

下面是原始的CSV格式文件，里面的数据表示了x、y、z的位置，最后还有能量值。x、y、z的坐标范围很大，下面是一个小片段——基本上这是在一个体积内进行的能量搜索。

-2.800000,-1.000000,5.470000,-0.26488315
-3.000000,1.000000,4.070000,-0.81185718
-2.800000,-1.000000,3.270000,1.29303723
-2.800000,-0.400000,4.870000,-0.51165026

不幸的是，想要在四个维度中绘图非常困难，所以我需要对这些数据进行整理。我希望能把这个体积的数据处理成一个在最低能量的z轴上的表面。对于小的数据集，这个过程很简单，在Excel中先按X排序，然后按Y排序，再按能量排序，最后删除所有高于最低能量的值。对于小数据集来说，这个方法很有效，但现在数据量大了，就变得很麻烦。

我尝试了各种方法，比如把CSV文件拆分并使用排序命令，但效果不太好。如果有人能给我一些建议，告诉我该怎么做，我会非常感激。

Answer 1

我觉得numpy的 lexsort 可以满足你的排序需求。

一般来说，你可以按照以下步骤进行：

1. 把csv文件读入numpy数组——你有没有试过用python的 csv包 或者numpy的 genfromtext() 函数？

2. 使用lexsort进行排序

3. 去掉不必要的行

补充：可以看看 这个相关的问题。

Answer 2

首先，使用 csv.reader 把数据读入一个包含元组的列表中。接着，根据 (x, y) 的值对这些数据进行排序。为了让数据更清晰，可以使用 命名元组 来标识每个字段。

然后，利用 itertools.groupby 来把相关的 (x, y) 数据点聚集在一起。对于每一组数据，使用 min 函数找出能量最低的那个数据点：

>>> import csv, collections, itertools

>>> raw_data = '''\
-2.800000,-1.000000,5.470000,-0.26488315
-3.000000,1.000000,4.070000,-0.81185718
-2.800000,-1.000000,3.270000,1.29303723
-2.800000,-0.400000,4.870000,-0.51165026
'''.splitlines()

>>> Sample = collections.namedtuple('Sample', ['x', 'y', 'z', 'energy'])
>>> data = [Sample(*row) for row in csv.reader(raw_data)]
>>> data.sort(key=lambda s: (s.x, s.y))
>>> for xy, group in itertools.groupby(data, key=lambda s: (s.x, s.y)):
        print min(group, key=lambda s: s.energy)


Sample(x='-2.800000', y='-0.400000', z='4.870000', energy='-0.51165026')
Sample(x='-2.800000', y='-1.000000', z='5.470000', energy='-0.26488315')
Sample(x='-3.000000', y='1.000000', z='4.070000', energy='-0.81185718')

Answer 3

这段代码实现了你在评论中对Raymond的回答所提到的功能——对于每一对 x, y，返回 z 值最低的那一行数据：

from operator import itemgetter
from itertools import groupby
from csv import reader


def min_z(iterable):
    # the data converted from strings to numbers
    floats = [[float(n) for n in row] for row in iterable]
    # the data sorted by x, y, z
    floats.sort(key=lambda (x, y, z, e): (x, y, z))
    # group the data by x, y
    grouped_floats = groupby(floats, key=itemgetter(slice(0, 2)))
    # return the first item from each group
    # because the data is sorted
    # the first item is the smallest z for the x, y group
    return [next(rowgroup) for xy, rowgroup in grouped_floats]


data = """-2.800000,-1.000000,5.470000,-0.26488315
-3.000000,1.000000,4.070000,-0.81185718
-2.800000,-1.000000,3.270000,1.29303723
-2.800000,-0.400000,4.870000,-0.51165026""".splitlines()


print min_z(reader(data))

输出结果是：

[[-3.0, 1.0, 4.07, -0.81185718], 
 [-2.8, -1.0, 3.27, 1.29303723], 
 [-2.8, -0.4, 4.87, -0.51165026]]