从CSV文件生成树

时间:2019-01-08 13:30:36

标签: python csv tree

我有以下格式的CSV文件:

Node, Count
A, 10000
A | A1,9000
A | A2, 500
A | A2 | A21, 400
A | A2 | A22, 300
A | A2 | A23, 100
A | A3, 800
A | A3 | A31, 500
A | A3 | A32, 350

我想根据这些数据构建一棵树,以便:

A has 3 child nodes - A1, A2, A3
A2 has 3 child nodes - A21,A22,A23
A3 has 2 child nodes - A31, A32

有什么建议吗?

1 个答案:

答案 0 :(得分:1)

您可以将“节点”列放入集字典中。 Dict将具有三个键A,A2和A3,并且这些键的值将是一个集合(以避免重复)

import csv

# dictionary with string keys and set() values
tree = {}

with open('csvfile.txt') as csv_file:
    csv_reader = csv.reader(csv_file, delimiter=',')
    line_count = 0
    for row in csv_reader:
        if line_count == 0:
            # skip header
            line_count +=1
        else:
            # grab the first column, split by |, strip whitespace
            nodes = [node.strip() for node in row[0].split('|')]
            for i in range(len(nodes) - 1):
                tree.setdefault(nodes[i], set()).add(nodes[i+1])
            line_count += 1

print(tree)
# output: {'A': {'A3', 'A2', 'A1'}, 'A2': {'A23', 'A22', 'A21'}, 'A3': {'A31', 'A32'}}