Data Structure and Algorithm - Find the Maximal Traffic between Cities

1. Problem Description:
This is a Google interview question for software engineer/developer from careercup. Here is the original thread,
"

You are given a graph with no cycles, each node representing different cities and there are stadiums for baseball games in all cities.

Each node contains a value representing the population of the city, and a list of neighbors. (feel free to extend data structure)

Every time there is a baseball game at a city, everyone from all the different cities will go to the city with the baseball game.

Return the maximum traffic between a city and its neighbours when there is a game at that city, for all cities. (Does not have to be sorted)

The total run-time after returning everything should be O(n).

Examples:

Input:
1   2
 \ /
  5
 / \
4   3
Output:
1 14
2 13
3 12
4 11
5 4

Input:
         38
         /
        8
        /
      7
     /
1   2
 \ / \
  5   15
 / \
4   3
Output:
1 82
2 53
3 80
4 79
5 70
7 46
15 68
8 38
38 45

- djway August 10, 2016 in United States | Report Duplicate | Flag 
".

2. Data Structure and Algorithm

The data structure and algorithm is based on graph.

Details please refer to my C++ post.

3. Python Implementation

    - Python 2.7

    - Microsoft Visual Studio Community 2015

# ********************************************************************************

# Implementation

# ********************************************************************************

# gNodeHelper.py
from exceptions import AttributeError
import sys

# dict() to take neighbour id as key and MaxTraffic as value
class gNode (object):
    def __init__(self, id, traffic):
        self.id = id
        self.traffic = traffic
        self.neighbours = dict()
        self.maxTraffic = 0
        self.sumOfNeighboursTraffic = 0

    def AppendNeighbour(self, neighbour):
        if neighbour is None:
            return
        self.neighbours[neighbour] = 0

    def AppendNeighbours(self, neighbours = None):
        if neighbours is None:
            return
        for neighbour in neighbours:
            self.AppendNeighbour(neighbour)

    # override the following 2 functions to enable gNode as hash key
    def __hash__(self):
        return hash(self.id)
    def __eq__(self, other):
        return self.id == other.id


class gNodeHelper (object):
    def __init__(self):
        pass

    @staticmethod
    def __CalculateNeighboursTraffic(root, parent):
        if root is None:
            return 0
        if len(root.neighbours) == 1 and parent is not None:
            return root.traffic

        traffic = 0
        for key in root.neighbours.keys():
            if parent is None or key.id != parent.id:
                root.neighbours[key] = gNodeHelper.__CalculateNeighboursTraffic(key, root)
                traffic += root.neighbours[key]
        return traffic + root.traffic

    @staticmethod
    def __CalculateParentTraffic(root, parent):
        if root is None:
            return
        if parent is not None:
            root.neighbours[parent] = parent.sumOfNeighboursTraffic + \
                                                      parent.traffic - parent.neighbours[root]
   
        for val in root.neighbours.values():
            root.sumOfNeighboursTraffic += val
   
        for key in root.neighbours.keys():
            if parent is None or key.id != parent.id:
                gNodeHelper.__CalculateParentTraffic(key, root)

    @staticmethod
    def __PopulateMaxTraffic(root, parent):
        if root is None:
            return
        for val in root.neighbours.values():
            if root.maxTraffic < val:
                root.maxTraffic = val
   
        for key in root.neighbours.keys():
            if parent is None or key != parent:
                gNodeHelper.__PopulateMaxTraffic(key, root)

    @staticmethod
    def GenerateMaxTraffic(root):
        try:
            gNodeHelper.__CalculateNeighboursTraffic(root, None)
            gNodeHelper.__CalculateParentTraffic(root, None)
            gNodeHelper.__PopulateMaxTraffic(root, None)
        except AttributeError() as e:
            print ("AttributeError: {0}".format(e))
        except:
            print ("Unexpected error: ", sys.exc_info()[0])

# ********************************************************************************
# Test
# ********************************************************************************
# Test_GenerateMaxTraffic.py
import unittest
from gNodeHelper import gNode
from gNodeHelper import gNodeHelper

class Test_GenerateMaxTraffic(unittest.TestCase):
    def test_A(self):
        node1 = gNode(1, 1)
        node2 = gNode(2, 2)
        node3 = gNode(3, 3)
        node4 = gNode(4, 4)
        node5 = gNode(5, 5)

        node1.AppendNeighbour(node5)
        node5.AppendNeighbours((node1, node2, node3, node4))
        node2.AppendNeighbour(node5)
        node3.AppendNeighbour(node5)
        node4.AppendNeighbour(node5)

        gNodeHelper.GenerateMaxTraffic(node1)

        self.assertEqual(node1.maxTraffic, 14)
        self.assertEqual(node2.maxTraffic, 13)
        self.assertEqual(node3.maxTraffic, 12)
        self.assertEqual(node4.maxTraffic, 11)
        self.assertEqual(node5.maxTraffic, 4)

    def test_B(self):
        node1 = gNode(1, 1)
        node2 = gNode(2, 2)
        node3 = gNode(3, 3)
        node4 = gNode(4, 4)
        node5 = gNode(5, 5)

        node1.AppendNeighbour(node5)
        node5.AppendNeighbours((node1, node2, node3, node4))
        node2.AppendNeighbour(node5)
        node3.AppendNeighbour(node5)
        node4.AppendNeighbour(node5)

        gNodeHelper.GenerateMaxTraffic(node5)

        self.assertEqual(node1.maxTraffic, 14)
        self.assertEqual(node2.maxTraffic, 13)
        self.assertEqual(node3.maxTraffic, 12)
        self.assertEqual(node4.maxTraffic, 11)
        self.assertEqual(node5.maxTraffic, 4)

    def test_C(self):
        node1 = gNode(1, 1)
        node2 = gNode(2, 2)
        node3 = gNode(3, 3)
        node4 = gNode(4, 4)
        node5 = gNode(5, 5)
        node7 = gNode(7, 7)
        node8 = gNode(8, 8)
        node15 = gNode(15, 15)
        node38 = gNode(38, 38)

        node1.AppendNeighbour(node5)
        node5.AppendNeighbours((node1, node2, node3, node4))
        node2.AppendNeighbours((node5, node7, node15))
        node3.AppendNeighbour(node5)
        node4.AppendNeighbour(node5)
        node7.AppendNeighbours((node2, node8))
        node8.AppendNeighbours((node7, node38))
        node15.AppendNeighbour(node2)
        node38.AppendNeighbour(node8)

        gNodeHelper.GenerateMaxTraffic(node1)

        self.assertEqual(node1.maxTraffic, 82)
        self.assertEqual(node2.maxTraffic, 53)
        self.assertEqual(node3.maxTraffic, 80)
        self.assertEqual(node4.maxTraffic, 79)
        self.assertEqual(node5.maxTraffic, 70)
        self.assertEqual(node7.maxTraffic, 46)
        self.assertEqual(node8.maxTraffic, 38)
        self.assertEqual(node15.maxTraffic, 68)
        self.assertEqual(node38.maxTraffic, 45)

    def test_D(self):
        node1 = gNode(1, 1)
        node2 = gNode(2, 2)
        node3 = gNode(3, 3)
        node4 = gNode(4, 4)
        node5 = gNode(5, 5)
        node7 = gNode(7, 7)
        node8 = gNode(8, 8)
        node15 = gNode(15, 15)
        node38 = gNode(38, 38)

        node1.AppendNeighbour(node5)
        node5.AppendNeighbours((node1, node2, node3, node4))
        node2.AppendNeighbours((node5, node7, node15))
        node3.AppendNeighbour(node5)
        node4.AppendNeighbour(node5)
        node7.AppendNeighbours((node2, node8))
        node8.AppendNeighbours((node7, node38))
        node15.AppendNeighbour(node2)
        node38.AppendNeighbour(node8)

        gNodeHelper.GenerateMaxTraffic(node5)

        self.assertEqual(node1.maxTraffic, 82)
        self.assertEqual(node2.maxTraffic, 53)
        self.assertEqual(node3.maxTraffic, 80)
        self.assertEqual(node4.maxTraffic, 79)
        self.assertEqual(node5.maxTraffic, 70)
        self.assertEqual(node7.maxTraffic, 46)
        self.assertEqual(node8.maxTraffic, 38)
        self.assertEqual(node15.maxTraffic, 68)
        self.assertEqual(node38.maxTraffic, 45)

if __name__ == '__main__':
    unittest.main()