How to access the dictionary values from classes in python?

Question

I am trying to use classes to create a dictionary using the following code:

import numpy as np
import collections


class Graph:
    def __init__(self):
        self.graph = collections.defaultdict(dict)

    def add_edge(self, u, v, weight = 1, directed = True):
        self.graph[u][v] = weight
        if not directed:
            self.graph[v][u] = weight

    def __str__(self):
        to_return = ''
        for vertex in self.graph:
            to_return += str(vertex) + ': '
            for edge in self.graph[vertex]:
                to_return +=  '(' + str(edge) + ', ' + str(self.graph[vertex][edge]) + ')'
                to_return += '   '

            to_return += '\n'
        return to_return

link_i = [1, 1, 3, 3, 3, 4, 4, 5, 5, 6]
link_j = [2, 3, 1, 2, 5, 5, 6, 6, 4, 4]

if __name__ == '__main__':
    g = Graph()
    for i in range(len(link_i)):
        g.add_edge(link_i[i],link_j[i])

    print(g)
    
for key in g.graph:
    print(g.graph[key])
 

This gives an output that are dictionaries

{2: 1, 3: 1}
{1: 1, 2: 1, 5: 1}
{5: 1, 6: 1}
{6: 1, 4: 1}
{4: 1}

How can I get the output as a list of arrays instead of a list of dictionaries?

Answer 1

You could provide methods to your class to wrap this functionality:

class Graph:
    # all the other stuff

    def verteces(self):
        return self.graph.keys()

    def neigbours(self, vertex):
        return list(self.graph[vertex].keys())

g = Graph()
# fill the graph

for v in g.verteces():
    print(g.neighbours(v))