Sub-tree with minimum color difference in a 2-coloured tree

A tree with N nodes and N-1 edges is given with 2 different colours for its nodes. 
Find the sub-tree with minimum colour difference i.e. abs(1-colour nodes – 2-colour nodes) is minimum. 

Example:

Input : 
Edges : 1 2
        1 3
        2 4
        3 5
Colours : 1 1 2 2 1 [1-based indexing where 
                    index denotes the node]
Output : 2
Explanation : The sub-tree {1-2} and {1-2-3-5}
have color difference of 2. Sub-tree {1-2} has two
1-colour nodes and zero 2-colour nodes. So, color 
difference is 2. Sub-tree {1-2-3-5} has three 1-colour
nodes and one 2-colour nodes. So color diff = 2.

Method 1 : The problem can be solved by checking every possible sub-tree from every node of the tree. This will take exponential time as we will check for sub-trees from every node.

Method 2 : (Efficient) If we observe, we are solving a portion of the tree several times. This produces recurring sub-problems. We can use Dynamic Programming approach to get the minimum color difference in one traversal. To make things simpler, we can have color values as 1 and -1. Now, if we have a sub-tree with both colored nodes equal, our sum of colors will be 0. To get the minimum difference, we should have maximum negative sum or maximum positive sum. 

• Case 1 When we need to have a sub-tree with maximum sum : We take a node if its value > 0, i.e. sum(parent) += max(0, sum(child))
• Case 2 When we need to have a sub-tree with minimum sum(or max negative sum) : We take a node if its value < 0, i.e. sum(parent) += min(0, sum(child))

To get the minimum sum, we can interchange the colors of nodes, i.e. -1 becomes 1 and vice-versa.

Below is the implementation : 

C++

// CPP code to find the sub-tree with minimum color
// difference in a 2-coloured tree
#include <bits/stdc++.h>
using namespace std;

// Tree traversal to compute minimum difference
void dfs(int node, int parent, vector<int> tree[],
                    int colour[], int answer[])
{
    // Initial min difference is the color of node
    answer[node] = colour[node];

    // Traversing its children
    for (auto u : tree[node]) {

        // Not traversing the parent
        if (u == parent)
            continue;

        dfs(u, node, tree, colour, answer);

        // If the child is adding positively to
        // difference, we include it in the answer
        // Otherwise, we leave the sub-tree and
        // include 0 (nothing) in the answer
        answer[node] += max(answer[u], 0);
    }
}

int maxDiff(vector<int> tree[], int colour[], int N)
{
       int answer[N + 1];
       memset(answer, 0, sizeof(answer));

    // DFS for colour difference : 1colour - 2colour
    dfs(1, 0, tree, colour, answer);

    // Minimum colour difference is maximum answer value
    int high = 0;
    for (int i = 1; i <= N; i++) {
        high = max(high, answer[i]);

        // Clearing the current value
        // to check for colour2 as well
        answer[i] = 0;
    }

    // Interchanging the colours
    for (int i = 1; i <= N; i++) {
        if (colour[i] == -1)
            colour[i] = 1;
        else
            colour[i] = -1;
    }

    // DFS for colour difference : 2colour - 1colour
    dfs(1, 0, tree, colour, answer);

    // Checking if colour2 makes the minimum colour
    // difference
    for (int i = 1; i < N; i++)
        high = max(high, answer[i]);

    return high;
}

// Driver code
int main()
{
    // Nodes
    int N = 5;

    // Adjacency list representation
    vector<int> tree[N + 1];

    // Edges
    tree[1].push_back(2);
    tree[2].push_back(1);

    tree[1].push_back(3);
    tree[3].push_back(1);

    tree[2].push_back(4);
    tree[4].push_back(2);

    tree[3].push_back(5);
    tree[5].push_back(3);

    // Index represent the colour of that node
    // There is no Node 0, so we start from
    // index 1 to N
    int colour[] = { 0, 1, 1, -1, -1, 1 };

    // Printing the result
    cout << maxDiff(tree,  colour,  N);

    return 0;
}

Java

// Java code to find the sub-tree
// with minimum color difference
// in a 2-coloured tree
import java.util.*;
class GFG
{

// Tree traversal to compute minimum difference
static void dfs(int node, int parent,
                Vector<Integer> tree[],
                int colour[], int answer[])
{
    // Initial min difference is
    // the color of node
    answer[node] = colour[node];

    // Traversing its children
    for (Integer u : tree[node])
    {

        // Not traversing the parent
        if (u == parent)
            continue;

        dfs(u, node, tree, colour, answer);

        // If the child is adding positively to
        // difference, we include it in the answer
        // Otherwise, we leave the sub-tree and
        // include 0 (nothing) in the answer
        answer[node] += Math.max(answer[u], 0);
    }
}

static int maxDiff(Vector<Integer> tree[],
                   int colour[], int N)
{
    int []answer = new int[N + 1];

    // DFS for colour difference : 1colour - 2colour
    dfs(1, 0, tree, colour, answer);

    // Minimum colour difference is
    // maximum answer value
    int high = 0;
    for (int i = 1; i <= N; i++)
    {
        high = Math.max(high, answer[i]);

        // Clearing the current value
        // to check for colour2 as well
        answer[i] = 0;
    }

    // Interchanging the colours
    for (int i = 1; i <= N; i++)
    {
        if (colour[i] == -1)
            colour[i] = 1;
        else
            colour[i] = -1;
    }

    // DFS for colour difference : 2colour - 1colour
    dfs(1, 0, tree, colour, answer);

    // Checking if colour2 makes the
    // minimum colour difference
    for (int i = 1; i < N; i++)
        high = Math.max(high, answer[i]);

    return high;
}

// Driver code
public static void main(String []args)
{

    // Nodes
    int N = 5;

    // Adjacency list representation
    Vector<Integer> tree[] = new Vector[N + 1];
    for(int i = 0; i < N + 1; i++)
        tree[i] = new Vector<Integer>();

    // Edges
    tree[1].add(2);
    tree[2].add(1);

    tree[1].add(3);
    tree[3].add(1);

    tree[2].add(4);
    tree[4].add(2);

    tree[3].add(5);
    tree[5].add(3);

    // Index represent the colour of that node
    // There is no Node 0, so we start from
    // index 1 to N
    int colour[] = { 0, 1, 1, -1, -1, 1 };

    // Printing the result
    System.out.println(maxDiff(tree, colour, N));
}
}

// This code is contributed by 29AjayKumar

Python3

# Python3 code to find the sub-tree
# with minimum color difference
# in a 2-coloured tree

# Tree traversal to compute minimum difference
def dfs(node, parent, tree, colour, answer):
    # Initial min difference is
    # the color of node
    answer[node] = colour[node]

    # Traversing its children
    for u in tree[node]:

        # Not traversing the parent
        if (u == parent):
            continue

        dfs(u, node, tree, colour, answer)

        # If the child is Adding positively to
        # difference, we include it in the answer
        # Otherwise, we leave the sub-tree and
        # include 0 (nothing) in the answer
        answer[node] += max(answer[u], 0)

def maxDiff(tree, colour, N):
    answer = [0 for _ in range(N+1)]

    # DFS for colour difference : 1colour - 2colour
    dfs(1, 0, tree, colour, answer)

    # Minimum colour difference is
    # maximum answer value
    high = 0
    for i in range(1, N+1):
        high = max(high, answer[i])

        # Clearing the current value
        # to check for colour2 as well
        answer[i] = 0

    # Interchanging the colours
    for i in range(1, N+1):
        if colour[i] == -1:
            colour[i] = 1
        else:
            colour[i] = -1

    # DFS for colour difference : 2colour - 1colour
    dfs(1, 0, tree, colour, answer)

    # Checking if colour2 makes the
    # minimum colour difference
    for i in range(1, N):
        high = max(high, answer[i])

    return high

# Driver code
# Nodes
N = 5

# Adjacency list representation
tree = [list() for _ in range(N+1)]

# Edges
tree[1].append(2)
tree[2].append(1)
tree[1].append(3)
tree[3].append(1)
tree[2].append(4)
tree[4].append(2)
tree[3].append(5)
tree[5].append(3)

# Index represent the colour of that node
# There is no Node 0, so we start from
# index 1 to N
colour = [0, 1, 1, -1, -1, 1]

# Printing the result
print(maxDiff(tree, colour, N))

# This code is contributed by nitibi9839.

C#

// C# code to find the sub-tree
// with minimum color difference
// in a 2-coloured tree
using System;
using System.Collections.Generic;

class GFG
{

// Tree traversal to compute minimum difference
static void dfs(int node, int parent,
                List<int> []tree,
                int []colour, int []answer)
{
    // Initial min difference is
    // the color of node
    answer[node] = colour[node];

    // Traversing its children
    foreach (int u in tree[node])
    {

        // Not traversing the parent
        if (u == parent)
            continue;

        dfs(u, node, tree, colour, answer);

        // If the child is Adding positively to
        // difference, we include it in the answer
        // Otherwise, we leave the sub-tree and
        // include 0 (nothing) in the answer
        answer[node] += Math.Max(answer[u], 0);
    }
}

static int maxDiff(List<int> []tree,
                         int []colour, int N)
{
    int []answer = new int[N + 1];

    // DFS for colour difference : 1colour - 2colour
    dfs(1, 0, tree, colour, answer);

    // Minimum colour difference is
    // maximum answer value
    int high = 0;
    for (int i = 1; i <= N; i++)
    {
        high = Math.Max(high, answer[i]);

        // Clearing the current value
        // to check for colour2 as well
        answer[i] = 0;
    }

    // Interchanging the colours
    for (int i = 1; i <= N; i++)
    {
        if (colour[i] == -1)
            colour[i] = 1;
        else
            colour[i] = -1;
    }

    // DFS for colour difference : 2colour - 1colour
    dfs(1, 0, tree, colour, answer);

    // Checking if colour2 makes the
    // minimum colour difference
    for (int i = 1; i < N; i++)
        high = Math.Max(high, answer[i]);

    return high;
}

// Driver code
public static void Main(String []args)
{

    // Nodes
    int N = 5;

    // Adjacency list representation
    List<int> []tree = new List<int>[N + 1];
    for(int i = 0; i < N + 1; i++)
        tree[i] = new List<int>();

    // Edges
    tree[1].Add(2);
    tree[2].Add(1);

    tree[1].Add(3);
    tree[3].Add(1);

    tree[2].Add(4);
    tree[4].Add(2);

    tree[3].Add(5);
    tree[5].Add(3);

    // Index represent the colour of that node
    // There is no Node 0, so we start from
    // index 1 to N
    int []colour = { 0, 1, 1, -1, -1, 1 };

    // Printing the result
    Console.WriteLine(maxDiff(tree, colour, N));
}
}

// This code is contributed by Rajput-Ji

Javascript

<script>

// JavaScript code to find the sub-tree
// with minimum color difference
// in a 2-coloured tree

// Tree traversal to compute minimum difference
function dfs(node, parent, tree, colour, answer)
{
    // Initial min difference is
    // the color of node
    answer[node] = colour[node];

    // Traversing its children
    for(var u of tree[node])
    {

        // Not traversing the parent
        if (u == parent)
            continue;

        dfs(u, node, tree, colour, answer);

        // If the child is Adding positively to
        // difference, we include it in the answer
        // Otherwise, we leave the sub-tree and
        // include 0 (nothing) in the answer
        answer[node] += Math.max(answer[u], 0);
    }
}

function maxDiff(tree, colour, N)
{
    var answer = Array(N+1).fill(0);

    // DFS for colour difference : 1colour - 2colour
    dfs(1, 0, tree, colour, answer);

    // Minimum colour difference is
    // maximum answer value
    var high = 0;
    for (var i = 1; i <= N; i++)
    {
        high = Math.max(high, answer[i]);

        // Clearing the current value
        // to check for colour2 as well
        answer[i] = 0;
    }

    // Interchanging the colours
    for (var i = 1; i <= N; i++)
    {
        if (colour[i] == -1)
            colour[i] = 1;
        else
            colour[i] = -1;
    }

    // DFS for colour difference : 2colour - 1colour
    dfs(1, 0, tree, colour, answer);

    // Checking if colour2 makes the
    // minimum colour difference
    for (var i = 1; i < N; i++)
        high = Math.max(high, answer[i]);

    return high;
}

// Driver code
// Nodes
var N = 5;
// Adjacency list representation
var tree = Array.from(Array(N+1), ()=>Array());

// Edges
tree[1].push(2);
tree[2].push(1);
tree[1].push(3);
tree[3].push(1);
tree[2].push(4);
tree[4].push(2);
tree[3].push(5);
tree[5].push(3);
// Index represent the colour of that node
// There is no Node 0, so we start from
// index 1 to N
var colour = [0, 1, 1, -1, -1, 1];
// Printing the result
document.write(maxDiff(tree, colour, N));


</script>

Output

2

Time complexity: O(n) where n is the number of nodes in the binary tree.
Auxiliary Space: O(n)