HackerRank Kingdom Division problem solution

In this HackerRank Kingdom Division problem solution King Arthur has a large kingdom that can be represented as a tree, where nodes correspond to cities and edges correspond to the roads between cities. The kingdom has a total of n cities numbered from 1 to n.

The King wants to divide his kingdom between his two children, Reggie and Betty, by giving each of them 0 or more cities; however, they don’t get along so he must divide the kingdom in such a way that they will not invade each other’s cities. The first sibling will invade the second sibling’s city if the second sibling has no other cities directly connected to it.

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.

HackerRank Kingdom Division 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()

Kingdom Division 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;
}