This notebook was prepared by [Donne Martin](https://github.com/donnemartin). Source and license info is on [GitHub](https://github.com/donnemartin/interactive-coding-challenges).

Solution Notebook

Problem: Implement a graph.

Constraints

  • Is the graph directed?
    • Yes
  • Do the edges have weights?
    • Yes

Test Cases

Input:

  • add_edge(source, destination, weight)
graph.add_edge(0, 1, 5)
graph.add_edge(0, 5, 2)
graph.add_edge(1, 2, 3)
graph.add_edge(2, 3, 4)
graph.add_edge(3, 4, 5)
graph.add_edge(3, 5, 6)
graph.add_edge(4, 0, 7)
graph.add_edge(5, 4, 8)
graph.add_edge(5, 2, 9)

Result:

  • source and destination nodes within graph are connected with specified weight.

Note:

  • We'll be using an OrderedDict to make the outputs more consistent and simplify our testing.
  • Graph will be used as a building block for more complex graph challenges.

Algorithm

Node

Node will keep track of its:

  • id
  • visit state
  • adjacent
    • key: node
    • value: weight

Graph

Graph will keep track of its:

  • nodes
    • key: node id
    • value: node

add_node

  • Create a node with the input id
  • Add the newly created node to the list of nodes

Complexity:

  • Time: O(1)
  • Space: O(1)

add_edge

  • If the source node is not in the list of nodes, add it
  • If the dest node is not in the list of nodes, add it
  • Add a connection from the source node to the dest node with the given edge weight

Complexity:

  • Time: O(1)
  • Space: O(1)

Code


In [1]:
%%writefile graph.py
from collections import OrderedDict
from enum import Enum  # Python 2 users: Run pip install enum34


class State(Enum):
    unvisited = 1
    visited = 2
    visiting = 3


class Node:

    def __init__(self, id):
        self.id = id
        self.visit_state = State.unvisited
        self.adjacent = OrderedDict()  # key = node, val = weight

    def __str__(self):
        return str(self.id)


class Graph:

    def __init__(self):
        self.nodes = OrderedDict()  # key = node id, val = node

    def add_node(self, id):
        node = Node(id)
        self.nodes[id] = node
        return node

    def add_edge(self, source, dest, weight=0):
        if source not in self.nodes:
            self.add_node(source)
        if dest not in self.nodes:
            self.add_node(dest)
        self.nodes[source].adjacent[self.nodes[dest]] = weight


Overwriting graph.py

In [2]:
%run graph.py

Unit Test


In [3]:
%%writefile test_graph.py
from nose.tools import assert_equal


class TestGraph(object):

    def test_graph(self):
        graph = Graph()
        for id in range(0, 6):
            graph.add_node(id)
        graph.add_edge(0, 1, 5)
        graph.add_edge(0, 5, 2)
        graph.add_edge(1, 2, 3)
        graph.add_edge(2, 3, 4)
        graph.add_edge(3, 4, 5)
        graph.add_edge(3, 5, 6)
        graph.add_edge(4, 0, 7)
        graph.add_edge(5, 4, 8)
        graph.add_edge(5, 2, 9)

        assert_equal(graph.nodes[0].adjacent[graph.nodes[1]], 5)
        assert_equal(graph.nodes[0].adjacent[graph.nodes[5]], 2)
        assert_equal(graph.nodes[1].adjacent[graph.nodes[2]], 3)
        assert_equal(graph.nodes[2].adjacent[graph.nodes[3]], 4)
        assert_equal(graph.nodes[3].adjacent[graph.nodes[4]], 5)
        assert_equal(graph.nodes[3].adjacent[graph.nodes[5]], 6)
        assert_equal(graph.nodes[4].adjacent[graph.nodes[0]], 7)
        assert_equal(graph.nodes[5].adjacent[graph.nodes[4]], 8)
        assert_equal(graph.nodes[5].adjacent[graph.nodes[2]], 9)

        print('Success: test_graph')


def main():
    test = TestGraph()
    test.test_graph()


if __name__ == '__main__':
    main()


Overwriting test_graph.py

In [4]:
%run -i test_graph.py


Success: test_graph