在Python中从CSV文件排序数据并生成表格输出

Question

假设我有一个这样的CSV文件（subjects.csv）

subjects,name1,name2,name3
Chemistry,Tom,Will,Rob
Biology,Megan,Sam,Tim
Physics,Tim,Will,Bob
Maths,Will,Tim,Joe

我想找出哪些学生是同班同学，主要关注Tim、Tom和Will。我该如何在Python中把他们配对呢？

比如说：

Tim和Will一起上了2门课。

Tom和Will一起上了1门课。

此外，我还想把这些信息放在一个表格里，像我下面写的那样，表格的两边都是名字，里面是每对学生共同上课的数量（名字可以按字母顺序排列）。我看过怎么为整个CSV文件生成表格，但我还是搞不懂怎么从头开始制作表格，同时又要从CSV文件中去掉一些列和行……

             Tim        Tom     Will

    Tim   0           0       0

    Tom   0           0       1

    Will     2           0       0

这超出了我个人的能力范围，但我还是想知道怎么做，并试着理解一下。

Answer 1

你可以创建一个字典，里面记录每个学生正在上哪些课：

>>> import csv
>>> import collections
>>> D = collections.defaultdict(set)
>>> with open('subjects.csv','rb') as f:
...     subject_reader = csv.reader(f)
...     header = subject_reader.next()
...     for row in subject_reader:
...         for name in row[1:]:
...             D[name].add(row[0])
... 
>>> import pprint
>>> pprint.pprint(dict(D))
{'Bob': set(['Physics']),
 'Joe': set(['Maths']),
 'Megan': set(['Biology']),
 'Rob': set(['Chemistry']),
 'Sam': set(['Biology']),
 'Tim': set(['Biology', 'Maths', 'Physics']),
 'Tom': set(['Chemistry']),
 'Will': set(['Chemistry', 'Maths', 'Physics'])}
>>> 

如果想检查有多少人一起上课，可以使用集合的交集方法：

>>> D['Tom'].intersection(D['Will'])
set(['Chemistry'])
>>> len(_)
1
>>> D['Tim'].intersection(D['Will'])
set(['Maths', 'Physics'])
>>> len(_)
2
>>> 

要打印出你例子中的表格，可以这样做：

>>> EXAMPLE_NAMES = ['Tom','Tim','Will']
>>> for y_name in EXAMPLE_NAMES:
...     print '{0:{width}}'.format(y_name,width=5),
...     for x_name in EXAMPLE_NAMES:
...         if y_name==x_name:
...             print '{0:{width}}'.format('-'*5, width=5),
...         else:
...             print '{0:{width}}'.format(len(D[y_name].intersection(D[x_name])), width=5),
...     print
... 
Tom   -----     0     1
Tim       0 -----     2
Will      1     2 -----

表格的标题可能看起来像这样：

    >>> for x_name in [' ']+EXAMPLE_NAMES:
    ...     print '{0:{width}}'.format(x_name, width=5),
    ... 
          Tom   Tim   Will 

正如约翰在评论中提到的，我把名字硬编码到一个列表里，以模拟你上面给的例子。要查看整个表格，你可以获取或遍历你创建的字典中的键，使用 .iterkeys() 或 .keys()：

>>> import csv
>>> import collections
>>> 
>>> my_d = collections.defaultdict(set)
>>> with open('subjects.csv','rb') as f:
...     subject_reader = csv.reader(f)
...     header = subject_reader.next()
...     for row in subject_reader:
...         for name in row[1:]:
...             my_d[name].add(row[0])
... 
>>> def display_header(D):
...     for x_name in [' ']+D.keys():
...         print '{0:{width}}'.format(x_name, width=5),
...     print
... 
>>> def display_body(D):
...     for y_name in D.iterkeys():
...         print '{0:{width}}'.format(y_name,width=5),
...         for x_name in D.iterkeys():
...             if y_name==x_name:
...                 print '{0:{width}}'.format('-'*5, width=5),
...             else:
...                 print '{0:{width}}'.format(len(D[y_name].intersection(D[x_name])), width=5),
...         print
... 
>>> def display_table(D):
...     display_header(D)
...     display_body(D)
... 
>>> display_table(my_d)
      Sam   Rob   Megan Will  Tim   Joe   Tom   Bob  
Sam   -----     0     1     0     1     0     0     0
Rob       0 -----     0     1     0     0     1     0
Megan     1     0 -----     0     1     0     0     0
Will      0     1     0 -----     2     1     1     1
Tim       1     0     1     2 -----     1     0     1
Joe       0     0     0     1     1 -----     0     0
Tom       0     1     0     1     0     0 -----     0
Bob       0     0     0     1     1     0     0 -----
>>>