In this HackerRank The Story of a Tree problem solution Alice and Bob play q games. Given the tree, Alice’s guesses, and the value of k for each game, find the probability that Alice will win the game and print it on a new line as a reduced fraction in the format p/q.
Problem solution in Python.
import sys
from math import gcd
from collections import defaultdict
sys.setrecursionlimit(10**9)
def dfs1():
'''keep 1 node as root and calculate how many arrows are pointing towards it'''
root = 1
seen = {root, }
stack = [root]
count = 0
while stack:
root = stack.pop()
for node in tree[root]:
if node not in seen:
seen.add(node)
count += (root, node) in guess # if root is parent of node
stack.append(node)
return count
def dfs2(cost, k):
'''now make every node as root and calculate how many nodes
are pointed towards it; If u is the root node for which
dfs1 calculated n (number of arrows pointed towards the root)
then for v (adjacent node of u), it would be n-1 as v is the
made the parent now in this step (only if there is a guess, if
there is no guess then it would be not changed)'''
root = 1
seen = {root,}
stack = [(root, cost)]
t_win = 0
while stack:
(root, cost) = stack.pop()
t_win += (cost >= k)
for node in tree[root]:
if node not in seen:
seen.add(node)
stack.append((node, cost - ((root, node) in guess) + ((node, root) in guess)))
return t_win
q = int(input().strip())
for a0 in range(q):
n = int(input().strip())
tree = defaultdict(list)
for a1 in range(n-1):
u,v = map(int, input().strip().split(' '))
tree[u].append(v)
tree[v].append(u)
g,k = map(int, input().strip().split(' '))
guess = {tuple(map(int, input().split())) for _ in range(g)}
cost = dfs1()
win = dfs2(cost, k)
g = gcd(win, n)
print("{0}/{1}".format(win//g, n//g))
Problem solution in Java.
import java.util.HashSet;
import java.util.LinkedList;
import java.util.Scanner;
public class Solution {
static int N, K;
static int num;
static LinkedList<Integer>[] adj;
static HashSet<Integer>[] outdegree;
static HashSet<Integer>[] indegree;
@SuppressWarnings("unchecked")
public static void main(String[] args) {
Scanner input = new Scanner(System.in);
int q = input.nextInt();
for (int Q = 0; Q < q; Q++) {
N = input.nextInt();
num = 0;
adj = new LinkedList[N]; outdegree = new HashSet[N]; indegree = new HashSet[N];
for (int i = 0; i < N; i++) {
adj[i] = new LinkedList<Integer>();
outdegree[i] = new HashSet<Integer>();
indegree[i] = new HashSet<Integer>();
}
for (int i = 0; i < N - 1; i++) {
int v = input.nextInt() - 1;
int w = input.nextInt() - 1;
adj[v].add(w);
adj[w].add(v);
}
int g = input.nextInt();
K = input.nextInt();
for (int G = 0; G < g; G++) {
int u = input.nextInt() - 1;
int v = input.nextInt() - 1;
indegree[u].add(v);
outdegree[v].add(u);
}
int v = 0;
walk(v, new boolean[N], init(v));
int gcd = GCD(N, num);
System.out.println((num / gcd) + "/" + (N / gcd));
}
}
static int GCD(int a, int b) {
if (a < b)
return GCD(b, a);
if (b == 0)
return a;
else
return GCD(b, a % b);
}
static void walk(int v, boolean[] visited, int amount) {
visited[v] = true;
if (amount >= K) num++;
for (int w : adj[v]) {
if (!visited[w]) {
int temp = amount;
if (indegree[v].contains(w)) temp--;
if (outdegree[v].contains(w)) temp++;
walk(w, visited, temp);
}
}
}
static int init(int v) {
return dfs(v, new boolean[N]);
}
static int dfs(int v, boolean[] visited) {
visited[v] = true;
int k = 0;
for (int w : adj[v]) {
if (!visited[w]) {
k += dfs(w, visited);
if (indegree[v].contains(w)) k++;
}
}
return k;
}
}
Problem solution in C++.
#include <vector>
#include <map>
#include <algorithm>
const int kN = 200000 + 5;
int n;
std::vector<int> edges[kN];
int dfn[kN], ed[kN], stmp, parent[kN];
void dfs(int u, int fa)
{
parent[u] = fa;
dfn[u] = stmp ++;
for (int v : edges[u]) if (v != fa) {
dfs(v, u);
}
ed[u] = stmp - 1;
}
int gcd(int a, int b)
{
return b == 0 ? a : gcd(b, a % b);
}
std::map<int, int> count;
struct SegmentTree {
bool in[kN << 2];
int t[kN << 2];
void modify(int L, int R, int dt, int l, int r, int rt)
{
if (R < l || r < L) return ;
in[rt] = true;
if (L <= l && r <= R) {
t[rt] += dt;
return ;
}
int mid = l + r >> 1;
modify(L, R, dt, l, mid, rt << 1);
modify(L, R, dt, mid + 1, r,rt << 1 | 1);
}
int calc(int val, int l, int r, int rt)
{
if (in[rt] == false) return 0;
int mid = l + r >> 1;
val += t[rt];
int total = r - l + 1;
if (in[rt] && l != r) {
total -= calc(val, l, mid, rt << 1);
total -= calc(val, mid + 1, r, rt << 1 | 1);
}
count[val] += total;
in[rt] = false;
t[rt] = 0;
return r - l + 1;
}
} sgt;
int main()
{
int cas;
scanf("%d", &cas);
while (cas--) {
scanf("%d", &n);
for (int i = 0; i < n; ++ i) edges[i].clear();
for (int i = 0; i < n - 1; ++ i) {
int a, b;
scanf("%d%d", &a, &b); a --; b --;
edges[a].emplace_back(b);
edges[b].emplace_back(a);
}
stmp = 0;
dfs(0, -1);
int g, k;
scanf("%d%d", &g, &k);
for (int i = 0; i < g; ++ i) {
int a, b;
scanf("%d%d", &a, &b); a --; b --;
if (parent[b] == a) {
sgt.modify(0, dfn[b] - 1, 1, 0, n - 1, 1);
sgt.modify(ed[b] + 1, n - 1, 1, 0, n - 1, 1);
} else {
sgt.modify(dfn[a], ed[a], 1, 0, n - 1, 1);
}
}
count.clear();
sgt.calc(0, 0, n - 1, 1);
int a = 0, b = n;
for (auto t : count) {
if (t.first >= k) {
a += t.second;
}
}
int r = gcd(a, b);
a /= r;
b /= r;
printf("%d/%dn", a, b);
}
}
Problem solution in C.
#include <assert.h>
#include <limits.h>
#include <math.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
struct Edge {
int x;
int y;
};
static int n;
static int nedges;
static int g;
static int k;
static struct Edge edge[200000];
static int first[100000];
static int last[100000];
static int parent[100000];
static int nwrong[100000];
static int nright[100000];
static int ngoodroots;
static int count;
static int compareedge(const void *a, const void *b)
{
const struct Edge *e1 = a;
const struct Edge *e2 = b;
if (e1->x != e2->x)
return e1->x - e2->x;
return e1->y - e2->y;
}
static void dfs(int node, int from)
{
int i;
parent[node] = from;
for (i = first[node]; i < last[node]; i++) {
assert(edge[i].x == node);
if (edge[i].y != from) {
dfs(edge[i].y, node);
}
}
}
static void dfs2(int node, int from)
{
int i;
assert(parent[node] == from);
//printf("visit %d (parent %d), count=%dn", node, from, count);
count += nwrong[node] - nright[node];
/*printf("add %d: nwrong[%d]=%d, nright[%d]=%dn",
nwrong[node]-nright[node],
node,nwrong[node],
node,nright[node]);
*/
if (count >= k) {
ngoodroots ++;
}
for (i = first[node]; i < last[node]; i++) {
assert(edge[i].x == node);
if (edge[i].y != from) {
dfs2(edge[i].y, node);
}
}
count -= nwrong[node] - nright[node];
}
static void one(void)
{
int i;
scanf("%d", &n);
memset(first, 0, n * sizeof *first);
memset(last, 0, n * sizeof *last);
memset(parent, 0, n * sizeof *parent);
memset(nwrong, 0, n * sizeof *nwrong);
memset(nright, 0, n * sizeof *nright);
ngoodroots = 0;
count = 0;
for (i = 0; i < n-1; i++) {
int x, y;
scanf("%d %d", &x, &y);
edge[2*i].x = x-1;
edge[2*i].y = y-1;
edge[2*i+1].x = y-1;
edge[2*i+1].y = x-1;
}
nedges = 2*(n-1);
qsort(edge, nedges, sizeof *edge, compareedge);
for (i = 0; i < nedges; i++) {
if (i == 0 || edge[i].x != edge[i-1].x)
first[edge[i].x] = i;
if (i == nedges-1 || edge[i].x != edge[i+1].x)
last[edge[i].x] = i+1;
}
dfs(0, 0);
scanf("%d %d", &g, &k);
for (i = 0; i < g; i++) {
int x, y;
scanf("%d %d", &x, &y);
if (parent[y-1] != x-1) {
nwrong[x-1] ++;
}
else {
nright[y-1] = 1;
count ++;
}
}
//printf("%d rightn", count);
dfs2(0, 0);
int a = ngoodroots;
int b = n;
for (i = n; i > 0; i--) {
if ((a % i)==0 && (b % i)==0) {
a /= i;
b /= i;
}
}
printf("%d/%dn", a, b);
}
int main(void)
{
int q;
scanf("%d", &q);
while (q--)
one();
return 0;
}
