Program to find maximum possible population of all the cities in python

Consider a country that is represented as a tree with N nodes and N-1 edges. Now each node represents a town, and each edge represents a road. We have two lists of numbers source and dest of size N-1. According to them the i-th road connects source[i] to dest[i]. And the roads are bidirectional. We also have another list of numbers called population of size N, where population[i] represents the population of the i-th town. We are trying to upgrade some number of towns into cities. But no two cities should be adjacent to each other and every node adjacent to a town should be a city (every road must connect a town and a city). So we have to find the maximum possible population of all the cities.

So, if the input is like source = [2, 2, 1, 1] dest = [1, 3, 4, 0] population = [6, 8, 4, 3, 5], then the output will be 15, as we can upgrade cities 0, 2, and 4 to get a population of 6 + 4 + 5 = 15.

To solve this, we will follow these steps −

• adj := make adjacency list of graph by using source and dest
• Define a function dfs() . This will take x, choose
• if x is seen, then
  • return 0
• mark x as seen
• ans := 0
• if choose is true, then
  • ans := ans + population[x]
• for each neighbor in adj[x], do
  • ans := ans + dfs(neighbor, inverse of choose)
• return ans
• From the main method do the following:
• x := dfs(0, True)
• return maximum of x and ((sum of population) - x)

Let us see the following implementation to get better understanding −

Example 

Live Demo

from collections import defaultdict
class Solution:
   def solve(self, source, dest, population):
      adj = defaultdict(list)
      for a, b in zip(source, dest):
         adj[a].append(b)
         adj[b].append(a)

      seen = set()

      def dfs(x, choose):
         if x in seen:
            return 0
         seen.add(x)
         ans = 0
         if choose:
            ans += population[x]
         for neighbor in adj[x]:
            ans += dfs(neighbor, not choose)
         return ans

      x = dfs(0, True)
      return max(x, sum(population) - x)
     
ob = Solution()
source = [2, 2, 1, 1]
dest = [1, 3, 4, 0]
population = [6, 8, 4, 3, 5]
print(ob.solve(source, dest, population))

Input

[2, 2, 1, 1], [1, 3, 4, 0], [6, 8, 4, 3, 5]

Output

15