HackerRank Kingdom Division problem solution

In this HackerRank Kingdom Division problem solution we have given a map of the kingdom’s n cities, find and print the number of ways King Arthur can divide it between his two children such that they will not invade each other. As this answer can be quite large, it must be modulo 10 to power 9 plus 7.

Problem solution in Python.

#!/bin/python3

import math
import os
import random
import re
import sys


def countColorings(root):
    v = {}
    for depth, node in d:
        if len(t[node]) == 0:
            v[(node, True)] = 1
            v[(node, False)] = 0
        else:
            cases = 1
            invalid = 1
            for child in t[node]:
                cases *= v[(child, True)] + v[(child, False)]
                invalid *= v[(child, False)]
            v[(node, True)] = cases
            v[(node, False)] = cases - invalid
    return v[(root, False)] * 2


def spanningTree(root):
    stack = [(root, 0)]
    while stack:
        node, depth = stack.pop()
        t[node] = set()
        d.append((depth, node))
        for adj in g[node]:
            if adj not in t:
                t[node].add(adj)
                stack.append((adj, depth + 1))


def kingdomDivision(n, roads):
    global g
    g = {}
    for road in roads:
        if road[0] in g:
            g[road[0]].add(road[1])
        else:
            g[road[0]] = {road[1]}
        if road[1] in g:
            g[road[1]].add(road[0])
        else:
            g[road[1]] = {road[0]}
    global t
    t = {}
    global d
    d = []
    spanningTree(n // 2)
    d.sort(reverse=True)
    return countColorings(n // 2) % (10 ** 9 + 7)


if __name__ == '__main__':
    fptr = open(os.environ['OUTPUT_PATH'], 'w')

    n = int(input())

    roads = []

    for _ in range(n-1):
        roads.append(list(map(int, input().rstrip().split())))

    result = kingdomDivision(n, roads)

    fptr.write(str(result) + 'n')

    fptr.close()

Problem solution in Java.

import java.io.*;
import java.util.*;
import java.text.*;
import java.math.*;
import java.util.regex.*;

public class Solution {
    //static final boolean RED = true;
    //static final boolean BLACK = false;
    static final int SAME = 1;
    static final int DIFF = 0;
    static final long MEM = 1000000007;
    
    static List<Node> readyNodes = new LinkedList<>();
    
    static class Node {
        int name;
        List<Node> neighbors = new ArrayList<>();
        Node parent = null;
        boolean isDone = false;
        
        public Node(int name) {
            this.name = name;
        }
        
        boolean checkReady() {
            for (Node child : neighbors) {
                if (!child.isDone) {
                    return false;
                }
            }
            return true;
        }
    }
    
    static void rootTree(Node child, Node parent) {
        child.parent = parent;
        child.neighbors.remove(parent);
        for (Node superChild : child.neighbors) {
            rootTree(superChild, child);
        }
        if (child.neighbors.isEmpty()) {
            readyNodes.add(child);
        }
    }
    
    /*
    static long howMany(Node child, boolean color, boolean parentColor) { // Need DP here instead
        if (child.neighbors.isEmpty()) {
            if (color == parentColor)
                return 1;
            else
                return 0;
        }
        
        long[] redChildren = new long[child.neighbors.size()];
        long[] blackChildren = new long[child.neighbors.size()];
        
        long ways = 1;
        for (int i = 0; i < redChildren.length; i++) {
            redChildren[i] = howMany(child.neighbors.get(i), RED, color);
            blackChildren[i] = howMany(child.neighbors.get(i), BLACK, color);
            
            ways = (ways * (redChildren[i] + blackChildren[i])) % MEM;
        }
        
        if (color != parentColor) {
            long sub = 1;
            if (color == RED) {
                for (int i = 0; i < blackChildren.length; i++) {
                    sub = (sub * blackChildren[i]) % MEM;
                    if (sub == 0)
                        break;
                }
            }
            else {
                for (int i = 0; i < redChildren.length; i++) {
                    sub = (sub * redChildren[i]) % MEM;
                    if (sub == 0)
                        break;
                }
            }
            ways = (ways - sub + MEM) % MEM;
        }
        
        return ways;
    }*/

    static long kingdomDivision(int n, Node[] cities) {
        rootTree(cities[0], null);
        
        long[][] ways = new long[n][2]; //ways[nodeName][Diff/Same]
        for (Node leaf : readyNodes) {
            ways[leaf.name][SAME] = 1;
            ways[leaf.name][DIFF] = 0;
            leaf.isDone = true;
        }
        
        while (!readyNodes.isEmpty()) {
            List<Node> nextNodes = new LinkedList<>();
            for (Node done : readyNodes) {
                if (done.parent != null && done.parent.checkReady())
                    nextNodes.add(done.parent);
            }
            readyNodes = nextNodes;
            for (Node node : readyNodes) {
                long w = 1;
                for (Node child : node.neighbors) {
                    w = (w * (ways[child.name][SAME] + ways[child.name][DIFF])) % MEM;
                }
                ways[node.name][SAME] = w;
            
                long sub = 1;
                for (Node child : node.neighbors) {
                    sub = (sub * (ways[child.name][DIFF])) % MEM;
                    if (sub == 0)
                        break;
                }
                ways[node.name][DIFF] = (w - sub + MEM) % MEM;
                node.isDone = true;
            }
        }
        
        return (2*ways[0][DIFF]) % MEM;
    }

    public static void main(String[] args) {
        Scanner in = new Scanner(System.in);
        int n = in.nextInt();
        Node[] cities = new Node[n];
        for (int i = 0; i < n; i++) {
            cities[i] = new Node(i);
        }
        for(int i = 0; i < n-1; i++){
            int city1 = in.nextInt() - 1;
            int city2 = in.nextInt() - 1;
            cities[city1].neighbors.add(cities[city2]);
            cities[city2].neighbors.add(cities[city1]);
        }
        long result = kingdomDivision(n, cities);
        System.out.println(result);
        in.close();
    }
}

Problem solution in C++.

#include <bits/stdc++.h>
#define MOD 1000000007LL
using namespace std;
int n,p1,p2;
vector <int> adj[100005];
typedef long long ll;
ll dp[100005][2][2];
ll f(int v, int p, int col, int pacol){
	if(dp[v][col][pacol]>=0ll) return dp[v][col][pacol];
	dp[v][col][pacol]=1ll;
	for(int x=0;x<adj[v].size();x++){
		if(adj[v][x]==p) continue;
		dp[v][col][pacol]*=f(adj[v][x],v,0,col)+f(adj[v][x],v,1,col);
		dp[v][col][pacol]%=MOD;
	}
	if(col!=pacol){
		ll temp=1ll;
		for(int x=0;x<adj[v].size();x++){
			if(adj[v][x]==p) continue;
			temp*=f(adj[v][x],v,pacol,col);
			temp%=MOD;
		}
		dp[v][col][pacol]=(dp[v][col][pacol]+MOD-temp)%MOD;
	}
	return dp[v][col][pacol];
}
int main(){
	memset(dp,-1,sizeof(dp));
	scanf("%d",&n);
	for(int x=0;x<n-1;x++){
		scanf("%d %d",&p1,&p2);
		adj[p1].push_back(p2);
		adj[p2].push_back(p1);
	}
	printf("%lldn",(f(1,0,0,1)+f(1,0,1,0))%MOD);
	return 0;
}

Problem solution in C.

#include <stdio.h>
#include <stdlib.h>
#define MOD 1000000007
typedef struct _lnode{
  int x;
  int w;
  struct _lnode *next;
} lnode;
void insert_edge(int x,int y,int w);
void dfs(int x,int y);
long long not_care[100000],safe[100000];
lnode *table[100000]={0};

int main(){
  int n,x,y,i;
  scanf("%d",&n);
  for(i=0;i<n-1;i++){
    scanf("%d%d",&x,&y);
    insert_edge(x-1,y-1,0);
  }
  dfs(0,-1);
  printf("%lld",safe[0]*2%MOD);
  return 0;
}
void insert_edge(int x,int y,int w){
  lnode *t=(lnode*)malloc(sizeof(lnode));
  t->x=y;
  t->w=w;
  t->next=table[x];
  table[x]=t;
  t=(lnode*)malloc(sizeof(lnode));
  t->x=x;
  t->w=w;
  t->next=table[y];
  table[y]=t;
  return;
}
void dfs(int x,int y){
  int f=0;
  long long not_safe=1;
  lnode *p;
  not_care[x]=1;
  for(p=table[x];p;p=p->next)
    if(p->x!=y){
      dfs(p->x,x);
      f=1;
      not_care[x]=(not_care[p->x]+safe[p->x])%MOD*not_care[x]%MOD;
      not_safe=not_safe*safe[p->x]%MOD;
    }
  if(!f)
    safe[x]=0;
  else
    safe[x]=(not_care[x]-not_safe+MOD)%MOD;
  return;
}