Построить список для каждой записи источника

Question

У меня есть таблица csv, содержащая два столбца: «Source_Polygon_ID» и «Neighbor_Polygon_ID». Я хочу создать список, содержащий серию «Neighbor_Polygon_ID», если элемент в списке имеет общий «Source_Polygon_ID».

Вот пример таблицы csv: таблица csv с двумя столбцами

Вот пример результата, который я хотел бы получить: образец обработанного результата в формате CSV

Я использовал pandas для чтения таблицы csv и использовал функцию groupby для сортировки «Neighbor_Polygon_ID» по их общему «Source_Polygon_ID». Однако это не помогает мне составить списки «Neighbor_Polygon_ID» для каждого «Source_Polygon_ID». Любые решения?

Answer 1

Вот небольшой фрагмент, который делает то, что вы хотите, только используя стандартные библиотеки Python. Я предположил, что отношение соседей является симметричным, так что если A является соседом B, то также B является соседом A (если это не так, просто удалите вторую строку в первом цикле for).

Я использую defaultdict, который является подклассом обычного словаря Python и ведет себя точно так же, за исключением случаев, когда вы пытаетесь получить доступ к ключу, которого нет в словаре. В этом случае он создает значение этого ключа в качестве значения по умолчанию, которое вы можете установить при создании словаря. В данном случае это пустой набор. Вот почему доступ к neighbours[row['source']] никогда не вызовет KeyError.

Остальная часть кода довольно проста. Я просто читаю входной файл и для каждого идентификатора строю набор соседних идентификаторов. После этого я записываю их в выходной файл.

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)])