蟻本をpythonで: Sec. 2-5 グラフ (の下準備)

下準備: もろもろのグラフの実装

pythonのdictionaryで実装 (隣接リストでの実装に対応)

標準入力から各edgeの情報を受け取って、dictionaryを更新していく。

有向グラフは {node1: node2} でnode1 →　node2のedgeを表現
無向グラフは有向グラフから {node1: node2, node2: node1}で表現
重み付きの場合は{node1: [node2, weight]}としてある。

import copy
def join_dic(a, b):
    for i in b:
        if (i in a):
            if type(a[i]) != list:
                a[i] = [a[i],b[i]]
                a[i] = list(set(a[i]))
                if len(a[i]) ==1:
                    a[i] = a[i][0]
            elif not (b[i] in a[i]):
                a[i].append(b[i])
        else:
            a[i] = b[i]

    return a

def d_to_ud(a):
    output = copy.deepcopy(a)
    for i in a:
        nodes = a[i]
        if type(nodes) ==int:
            join_dic(output,{nodes:i})
        else:
            for j in nodes:
                join_dic(output,{j:i})
    return output


graph = {}
while 1:
    edge = input()
    if edge  == "":
        break
    edge = list(map(int,edge.split()))
    dic = {edge[0]:edge[1]}
    graph = join_dic(graph,dic)

print(graph)

print(d_to_ud(graph))

def join_weighted_graph(a, b):
    ## a = {from:[to, weight]}

    for i in b:

        if i in a:
          if type(a[i][0]) == list:
            a[i].append(b[i])
          else:
            a[i] = [a[i],b[i]]
        else:
          a[i] = b[i]

    return a


def d_to_ud_weighted(a):
    output = copy.deepcopy(a)
    for i in a:

        if type(a[i][0]) == list:

            for j in a[i]:

               node = j[0]
               weight = j[1]

               new_edge = {node:[i,weight]}

               join_weighted_graph(output,new_edge)

        else:
            join_weighted_graph(output,{a[i][0]:[i,a[i][1]]})
    return output


graph = {}    
while 1:
    edge = input()
    if edge == "":
        break
    edge = list(map(int,edge.split()))

    graph = join_weighted_graph(graph, {edge[0]:[edge[1],edge[2]]})

print(graph)

print(d_to_ud_weighted(graph))

追記)

前半:
入力

0 1
1 2
2 0

出力

{0: 1, 1: 2, 2: 0}
{0: [1, 2], 1: [0, 2], 2: [0, 1]}

後半:

入力

1 2 3
1 4 2
2 4 -1
2 3 4
3 4 5

出力

{1: [[2, 3], [4, 2]], 2: [[4, -1], [3, 4]], 3: [4, 5]}
{1: [[2, 3], [4, 2]], 2: [[4, -1], [3, 4], [1, 3]], 3: [[4, 5], [2, 4]], 4: [[1, 2], [2, -1], [3, 5]]}

追記2) autopep8を使ってpep8準拠に

import copy


def join_dic(a, b):
    for i in b:
        if (i in a):
            if type(a[i]) != list:
                a[i] = [a[i], b[i]]
                a[i] = list(set(a[i]))
                if len(a[i]) == 1:
                    a[i] = a[i][0]
            elif not (b[i] in a[i]):
                a[i].append(b[i])
        else:
            a[i] = b[i]

    return a


def d_to_ud(a):
    output = copy.deepcopy(a)
    for i in a:
        nodes = a[i]
        if type(nodes) == int:
            join_dic(output, {nodes: i})
        else:
            for j in nodes:
                join_dic(output, {j: i})
    return output


graph = {}
while 1:
    edge = input()
    if edge == "":
        break
    edge = list(map(int, edge.split()))
    dic = {edge[0]: edge[1]}
    graph = join_dic(graph, dic)

print(graph)

print(d_to_ud(graph))


def join_weighted_graph(a, b):
    # a = {from:[to, weight]}

    for i in b:

        if i in a:
            if type(a[i][0]) == list:
                a[i].append(b[i])
            else:
                a[i] = [a[i], b[i]]
        else:
            a[i] = b[i]

    return a


def d_to_ud_weighted(a):
    output = copy.deepcopy(a)
    for i in a:

        if type(a[i][0]) == list:

            for j in a[i]:

                node = j[0]
                weight = j[1]

                new_edge = {node: [i, weight]}

                join_weighted_graph(output, new_edge)

        else:
            join_weighted_graph(output, {a[i][0]: [i, a[i][1]]})
    return output


graph = {}
while 1:
    edge = input()
    if edge == "":
        break
    edge = list(map(int, edge.split()))

    graph = join_weighted_graph(graph, {edge[0]: [edge[1], edge[2]]})

print(graph)

print(d_to_ud_weighted(graph))