为每个源记录构造一个列表

Question

我这里有一个包含两列的 csv 表：“Source_Polygon_ID”和“Neighbor_Polygon_ID”。我想制作一个包含一系列“Neighbor_Polygon_ID”的列表，只要列表中的项目共享共同的“Source_Polygon_ID”。

以下是 csv 表的示例： csv table with two columns

这是我想要得到的示例结果： sample of processed result in csv format

我使用 pandas 来读取 csv 表，并使用 groupby 函数按其常见的“Source_Polygon_ID”对“Neighbor_Polygon_ID”进行排序。但是，这并不能帮助我为每个“Source_Polygon_ID”构建“Neighbor_Polygon_ID”列表。有什么解决办法吗？

Answer 1

这是一个小型的 sn-p，它只使用标准的 python 库来做你想做的事情。我假设邻居关系是对称的，所以如果 A 是 B 的邻居，那么 B 也是 A 的邻居（如果不是这种情况，那么只需删除第一个 for 循环中的第二行）。

我正在使用defaultdict，它是普通 python 字典的子类，其行为与它完全相同，除非您尝试访问不在字典中的键。在这种情况下，它将这个键的值实例化为默认值，您可以在创建字典时设置该值。在这种情况下，它是一个空集。这就是为什么访问neighbours[row['source']] 永远不会抛出KeyError。

其余的代码非常简单，我只是读取输入文件并为每个 id 构建一组相邻的 id。之后，我将这些写入输出文件。

import csv
from collections import defaultdict

neighbours = defaultdict(set)

with open('input.csv') as f:
    reader = csv.DictReader(f)
    for row in reader:
        neighbours[row['source']].add(row['neighbour'])
        neighbours[row['neighbour']].add(row['source'])

with open('output.csv', 'w') as f:
    writer = csv.writer(f)
    # Write the header
    writer.writerow(['source', 'neighbours'])

    # Write the data
    for key, values in neighbours.items():
        writer.writerow([key, ','.join(values)])