Balanced Forest

                            Solution in C :

In C ++ :



#include <cstdio>
#include <iostream>
#include <cstring>
#include <algorithm>
#include <cmath>
#include <vector>
#include <map>
#include <set>
#include <string>
#include <cstdlib>
#include <ctime>
#include <deque>
#include <unordered_set>
using namespace std;

int q;
map <long long, int> Map1, Map2;
long long ctot;
int c[110000];
vector <int> ve[110000];
long long ans;
int n;
long long sum[1100000];

void dfs1(int x, int f) {
	sum[x] = c[x];
	for (int i = 0; i < (int) ve[x].size(); i++)
		if (ve[x][i] != f) {
			dfs1(ve[x][i], x);
			sum[x] += sum[ve[x][i]];
		}
	Map1[sum[x]] += 1;
}

void test(long long x) {
	long long y = ctot - 2 * x;
	if (y > 0 && y <= x)
		ans = min(ans, x - y);
}

void dfs2(int x, int f) {
	

	
	if (Map2[2 * sum[x]])
		test(sum[x]);
	if (Map2[ctot - sum[x]])
		test(sum[x]);
	if ((ctot - sum[x]) % 2 == 0 && Map2[ctot - (ctot - sum[x]) / 2])
		test((ctot - sum[x]) / 2);

	Map2[sum[x]] += 1;

	if (Map1[sum[x]] > Map2[sum[x]])
		test(sum[x]);

	if (ctot - 2 * sum[x] >= sum[x] && Map1[ctot - 2 * sum[x]] > Map2[ctot - 2 * sum[x]])
		test(sum[x]);

	if ((ctot - sum[x]) % 2 == 0 && (ctot - sum[x]) / 2 >= sum[x] && Map1[(ctot - sum[x]) / 2] > Map2[(ctot - sum[x]) / 2])
		test((ctot - sum[x]) / 2);

	if (sum[x] * 2 == ctot)
		ans = min(ans, sum[x]);
	
	for (int i = 0; i < (int) ve[x].size(); i++)
		if (ve[x][i] != f) {
			dfs2(ve[x][i], x);
		}

	Map2[sum[x]] -= 1;
}

int main() {
	scanf("%d", &q);
	while (q--) {
		Map1.clear();
		Map2.clear();
		ans = 1e18;
		scanf("%d", &n);
		ctot = 0;
		for (int i = 1; i <= n; i++) {
			scanf("%d", &c[i]);
			ctot += c[i];
		}
		for (int i = 1; i <= n; i++)
			ve[i].clear();
		for (int i = 1; i < n; i++) {
			int x, y;
			scanf("%d%d", &x, &y);
			ve[x].push_back(y);
			ve[y].push_back(x);
		}
		dfs1(1, 0);
		dfs2(1, 0);
		if (ans == 1e18)
			printf("-1\n");
		else
			printf("%lld\n", ans);
	}
}







In Java :






import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.PrintWriter;
import java.util.Arrays;
import java.util.BitSet;
import java.util.Comparator;
import java.util.InputMismatchException;

public class E {
	InputStream is;
	PrintWriter out;
	String INPUT = "";
	
	void solve()
	{
		for(int T = ni();T > 0;T--){
			int n = ni();
			int[] a = na(n);
			int[] from = new int[n - 1];
			int[] to = new int[n - 1];
			for (int i = 0; i < n - 1; i++) {
				from[i] = ni() - 1;
				to[i] = ni() - 1;
			}
			int[][] g = packU(n, from, to);
			int[][] pars = parents3(g, 0);
			int[] par = pars[0], ord = pars[1], dep = pars[2];
			int[][] rs = makeRights(g, par, 0);
			int[] iord = rs[1], right = rs[2];
			int[][] spar = logstepParents(par);
			
			long[] des = new long[n];
			for(int i = n-1;i >= 0;i--){
				int cur = ord[i];
				des[cur] += a[cur];
				if(i > 0)des[par[cur]] += des[cur];
			}
			
			long W = des[0];
			
			long ret = Long.MAX_VALUE;
			// t=u
			{
				long[] dd = Arrays.copyOf(des, n);
				Arrays.sort(dd);
				for(int i = 0;i < n-1;i++){
					if(dd[i] == dd[i+1] && W-2*dd[i] <= dd[i] && W-2*dd[i] >= 0){
						ret = Math.min(ret, dd[i]-(W-2*dd[i]));
					}
				}
			}
			
			// s=t
			{
				long[][] poss = new long[n][];
				for(int i = 0;i < n;i++){
					poss[i] = new long[]{des[i], iord[i]};
				}
				Arrays.sort(poss, new Comparator<long[]>() {
					public int compare(long[] a, long[] b) {
						if(a[0] != b[0])return Long.compare(a[0], b[0]);
						return Long.compare(a[1], b[1]);
					}
				});
				long[] posv = new long[n];
				for(int i = 0;i < n;i++)posv[i] = poss[i][0];
				for(int i = 0;i < n;i++){
					long t = des[i];
					long u = W-2L*t;
					if(u >= 0 && t >= u){
						int lb = lowerBound(posv, u);
						int ub = lowerBound(posv, u+1);
						if(lb < ub){
							if((int)poss[lb][1] < iord[i]){
								ret = Math.min(ret, t-u);
							}
							if((int)poss[ub-1][1] > right[iord[i]]){
								ret = Math.min(ret, t-u);
							}
						}
					}
				}
			}
			
			// ireko
			// t=u
			{
				for(int i = 0;i < n;i++){
					long u = des[i];
					long s = W-u*2;
					if(s >= 0 && s <= u){
						int cur = i;
						for(int h = spar.length-1;h >= 0;h--){
							int anc = spar[h][cur];
							if(anc == -1)continue;
							if(des[anc] <= 2*u){
								cur = anc;
							}
						}
						if(des[cur] == 2*u){
							ret = Math.min(ret, u-s);
						}
					}
				}
			}
			
			// s=u
			{
				for(int i = 0;i < n;i++){
					long u = des[i];
					long t = W-u*2;
					if(t >= 0 && t <= u){
						int cur = i;
						for(int h = spar.length-1;h >= 0;h--){
							int anc = spar[h][cur];
							if(anc == -1)continue;
							if(des[anc] <= t+u){
								cur = anc;
							}
						}
						if(des[cur] == t+u){
							ret = Math.min(ret, u-t);
						}
					}
				}
			}
			// s=t
			{
				for(int i = 0;i < n;i++){
					long u = des[i];
					if((W-u)%2 != 0)continue;
					long t = (W-u)/2;
					if(t >= 0 && t >= u){
						int cur = i;
						for(int h = spar.length-1;h >= 0;h--){
							int anc = spar[h][cur];
							if(anc == -1)continue;
							if(des[anc] <= t+u){
								cur = anc;
							}
						}
						if(des[cur] == t+u){
							ret = Math.min(ret, t-u);
						}
					}
				}
			}
			if(ret == Long.MAX_VALUE){
				out.println(-1);
			}else{
				out.println(ret);
			}
		}
	}
	
	public static int lowerBound(long[] a, long v)
	{
		int low = -1, high = a.length;
		while(high-low > 1){
			int h = high+low>>>1;
			if(a[h] >= v){
				high = h;
			}else{
				low = h;
			}
		}
		return high;
	}

	
	public static int[] sortByPreorder(int[][] g, int root){
		int n = g.length;
		int[] stack = new int[n];
		int[] ord = new int[n];
		BitSet ved = new BitSet();
		stack[0] = root;
		int p = 1;
		int r = 0;
		ved.set(root);
		while(p > 0){
			int cur = stack[p-1];
			ord[r++] = cur;
			p--;
			for(int e : g[cur]){
				if(!ved.get(e)){
					stack[p++] = e;
					ved.set(e);
				}
			}
		}
		return ord;
	}
	
	public static int[][] makeRights(int[][] g, int[] par, int root)
	{
		int n = g.length;
		int[] ord = sortByPreorder(g, root);
		int[] iord = new int[n];
		for(int i = 0;i < n;i++)iord[ord[i]] = i;
		
		int[] right = new int[n];
		for(int i = n-1;i >= 0;i--){
			int v = i;
			for(int e : g[ord[i]]){
				if(e != par[ord[i]]){
					v = Math.max(v, right[iord[e]]);
				}
			}
			right[i] = v;
		}
		return new int[][]{ord, iord, right};
	}

	
	public static int lca2(int a, int b, int[][] spar, int[] depth) {
		if (depth[a] < depth[b]) {
			b = ancestor(b, depth[b] - depth[a], spar);
		} else if (depth[a] > depth[b]) {
			a = ancestor(a, depth[a] - depth[b], spar);
		}

		if (a == b)
			return a;
		int sa = a, sb = b;
		for (int low = 0, high = depth[a], t = Integer.highestOneBit(high), k = Integer
				.numberOfTrailingZeros(t); t > 0; t >>>= 1, k--) {
			if ((low ^ high) >= t) {
				if (spar[k][sa] != spar[k][sb]) {
					low |= t;
					sa = spar[k][sa];
					sb = spar[k][sb];
				} else {
					high = low | t - 1;
				}
			}
		}
		return spar[0][sa];
	}

	protected static int ancestor(int a, int m, int[][] spar) {
		for (int i = 0; m > 0 && a != -1; m >>>= 1, i++) {
			if ((m & 1) == 1)
				a = spar[i][a];
		}
		return a;
	}

	public static int[][] logstepParents(int[] par) {
		int n = par.length;
		int m = Integer.numberOfTrailingZeros(Integer.highestOneBit(n - 1)) + 1;
		int[][] pars = new int[m][n];
		pars[0] = par;
		for (int j = 1; j < m; j++) {
			for (int i = 0; i < n; i++) {
				pars[j][i] = pars[j - 1][i] == -1 ? -1 : pars[j - 1][pars[j - 1][i]];
			}
		}
		return pars;
	}


	public static int[][] parents3(int[][] g, int root) {
		int n = g.length;
		int[] par = new int[n];
		Arrays.fill(par, -1);

		int[] depth = new int[n];
		depth[0] = 0;

		int[] q = new int[n];
		q[0] = root;
		for (int p = 0, r = 1; p < r; p++) {
			int cur = q[p];
			for (int nex : g[cur]) {
				if (par[cur] != nex) {
					q[r++] = nex;
					par[nex] = cur;
					depth[nex] = depth[cur] + 1;
				}
			}
		}
		return new int[][] { par, q, depth };
	}

	static int[][] packU(int n, int[] from, int[] to) {
		int[][] g = new int[n][];
		int[] p = new int[n];
		for (int f : from)
			p[f]++;
		for (int t : to)
			p[t]++;
		for (int i = 0; i < n; i++)
			g[i] = new int[p[i]];
		for (int i = 0; i < from.length; i++) {
			g[from[i]][--p[from[i]]] = to[i];
			g[to[i]][--p[to[i]]] = from[i];
		}
		return g;
	}

	
	void run() throws Exception
	{
		is = INPUT.isEmpty() ? System.in : new ByteArrayInputStream(INPUT.getBytes());
		out = new PrintWriter(System.out);
		
		long s = System.currentTimeMillis();
		solve();
		out.flush();
		if(!INPUT.isEmpty())tr(System.currentTimeMillis()-s+"ms");
	}
	
	public static void main(String[] args) throws Exception { new E().run(); }
	
	private byte[] inbuf = new byte[1024];
	private int lenbuf = 0, ptrbuf = 0;
	
	private int readByte()
	{
		if(lenbuf == -1)throw new InputMismatchException();
		if(ptrbuf >= lenbuf){
			ptrbuf = 0;
			try { lenbuf = is.read(inbuf); } catch (IOException e) { throw new InputMismatchException(); }
			if(lenbuf <= 0)return -1;
		}
		return inbuf[ptrbuf++];
	}
	
	private boolean isSpaceChar(int c) { return !(c >= 33 && c <= 126); }
	private int skip() { int b; while((b = readByte()) != -1 && isSpaceChar(b)); return b; }
	
	private double nd() { return Double.parseDouble(ns()); }
	private char nc() { return (char)skip(); }
	
	private String ns()
	{
		int b = skip();
		StringBuilder sb = new StringBuilder();
		while(!(isSpaceChar(b))){ // when nextLine, (isSpaceChar(b) && b != ' ')
			sb.appendCodePoint(b);
			b = readByte();
		}
		return sb.toString();
	}
	
	private char[] ns(int n)
	{
		char[] buf = new char[n];
		int b = skip(), p = 0;
		while(p < n && !(isSpaceChar(b))){
			buf[p++] = (char)b;
			b = readByte();
		}
		return n == p ? buf : Arrays.copyOf(buf, p);
	}
	
	private char[][] nm(int n, int m)
	{
		char[][] map = new char[n][];
		for(int i = 0;i < n;i++)map[i] = ns(m);
		return map;
	}
	
	private int[] na(int n)
	{
		int[] a = new int[n];
		for(int i = 0;i < n;i++)a[i] = ni();
		return a;
	}
	
	private int ni()
	{
		int num = 0, b;
		boolean minus = false;
		while((b = readByte()) != -1 && !((b >= '0' && b <= '9') || b == '-'));
		if(b == '-'){
			minus = true;
			b = readByte();
		}
		
		while(true){
			if(b >= '0' && b <= '9'){
				num = num * 10 + (b - '0');
			}else{
				return minus ? -num : num;
			}
			b = readByte();
		}
	}
	
	private long nl()
	{
		long num = 0;
		int b;
		boolean minus = false;
		while((b = readByte()) != -1 && !((b >= '0' && b <= '9') || b == '-'));
		if(b == '-'){
			minus = true;
			b = readByte();
		}
		
		while(true){
			if(b >= '0' && b <= '9'){
				num = num * 10 + (b - '0');
			}else{
				return minus ? -num : num;
			}
			b = readByte();
		}
	}
	
	private static void tr(Object... o) 
{ System.out.println(Arrays.deepToString(o)); }
}









In C :





#include <stdio.h>
#include <stdlib.h>
#define HASH_SIZE 123455
typedef struct _lnode{
  int x;
  int w;
  struct _lnode *next;
} lnode;
typedef struct _node{
  long long x;
  int c;
  long long ans;
  struct _node *next;
} node;
void clean_table();
void insert_edge(int x,int y,int w);
void dfs0(int x,int y);
void dfs1(int x,int y);
void insert(long long x,long long ans);
void removee(long long x,long long ans);
void search(long long x);
void freehash();
int a[50000];
long long sub[50000],min,sum;
lnode *table[50000]={0};
node *hash[HASH_SIZE]={0};

int main(){
  int T,n,x,y,i;
  scanf("%d",&T);
  while(T--){
    scanf("%d",&n);
    for(i=sum=0;i<n;i++){
      scanf("%d",a+i);
      sum+=a[i];
    }
    for(i=0;i<n-1;i++){
      scanf("%d%d",&x,&y);
      insert_edge(x-1,y-1,1);
    }
    dfs0(0,-1);
    min=-1;
    dfs1(0,-1);
    printf("%lld\n",min);
    clean_table();
    freehash();
  }
  return 0;
}
void clean_table(){
  int i;
  lnode *p,*pp;
  for(i=0;i<50000;i++)
    if(table[i]){
      p=table[i];
      while(p){
        pp=p->next;
        free(p);
        p=pp;
      }
      table[i]=NULL;
    }
  return;
}
void insert_edge(int x,int y,int w){
  lnode *t=malloc(sizeof(lnode));
  t->x=y;
  t->w=w;
  t->next=table[x];
  table[x]=t;
  t=malloc(sizeof(lnode));
  t->x=x;
  t->w=w;
  t->next=table[y];
  table[y]=t;
  return;
}
void dfs0(int x,int y){
  lnode *p;
  sub[x]=a[x];
  for(p=table[x];p;p=p->next)
    if(p->x!=y){
      dfs0(p->x,x);
      sub[x]+=sub[p->x];
    }
  return;
}
void dfs1(int x,int y){
  lnode *p;
  long long down,up;
  search(sub[x]);
  down=sub[x];
  up=sum-sub[x];
  if(down==up && min==-1)
    min=down;
  if(down%2==0 && down/2*3>=sum)
    insert(down/2,down/2*3-sum);
  if(up<down && up*3>=sum){
    insert(up,up*3-sum);
    insert(down-up,up*3-sum);
  }
  for(p=table[x];p;p=p->next)
    if(p->x!=y)
      dfs1(p->x,x);
  if(down%2==0 && down/2*3>=sum)
    removee(down/2,down/2*3-sum);
  if(up<down && up*3>=sum){
    removee(up,up*3-sum);
    removee(down-up,up*3-sum);
  }
  if(up%2==0 && up/2*3>=sum)
    insert(up/2,up/2*3-sum);
  if(down<up && down*3>=sum){
    insert(down,down*3-sum);
    insert(up-down,down*3-sum);
  }
  return;
}
void insert(long long x,long long ans){
  int bucket=x%HASH_SIZE;
  node *t=hash[bucket];
  while(t){
    if(t->x==x && t->ans==ans){
      t->c++;
      return;
    }
    t=t->next;
  }
  t=(node*)malloc(sizeof(node));
  t->x=x;
  t->ans=ans;
  t->c=1;
  t->next=hash[bucket];
  hash[bucket]=t;
  return;
}
void removee(long long x,long long ans){
  int bucket=x%HASH_SIZE;
  node *t=hash[bucket],*p=NULL;
  while(t){
    if(t->x==x && t->ans==ans){
      t->c--;
      if(!t->c){
        if(!p){
          hash[bucket]=t->next;
          free(t);
        }
        else{
          p->next=t->next;
          free(t);
        }
      }
      return;
    }
    p=t;
    t=t->next;
  }
  return;
}
void search(long long x){
  int bucket=x%HASH_SIZE;
  node *t=hash[bucket];
  while(t){
    if(t->x==x)
      if(min==-1 || t->ans<min)
        min=t->ans;
    t=t->next;
  }
  return;
}
void freehash(){
  int i;
  node *t,*p;
  for(i=0;i<HASH_SIZE;i++){
    t=hash[i];
    while(t){
      p=t->next;
      free(t);
      t=p;
    }
    hash[i]=NULL;
  }
  return;
}









In python3 :







from random import randrange
import sys
sys.setrecursionlimit(10**5)

def r():
    return list(map(int, input().split()))

def bin_search(arr, pred, lo=0):
    hi = len(arr)
    while lo < hi:
        mid = (lo+hi)//2
        if pred(arr[mid]):
            hi = mid
        else: 
            lo = mid + 1
    return lo if lo < len(arr) else None

class UndirectedGraph(object):
    def __init__(self, size):
        self.size = size
        self.M = 0
        self.edges = [set([]) for _ in range(self.size)]

    def add_edge(self, u, v):
        """Adds edge to graph if it doesn't already exist."""
        if v not in self.edges[u]:
            self.edges[u].add(v)
            self.edges[v].add(u)
            self.M += 1

    def neighbors(self, u):
        for v in self.edges[u]:
            yield v

class WeightedTree(UndirectedGraph):
    def __init__(self, size, weights):
        super().__init__(size)
        self.weights = weights
        self.labels = [None]*self.size
        self.cum_weights = [None]*self.size
        self.inverted = [False]*self.size
        self.nextnode = 0
        self.num_children = [None]*self.size
        self.root = 0

    def initialize(self):
        self.set_labels(self.root)
        self.sorted_nodes = sorted(list(range(self.size)), key = lambda x : self.cum_weights[x])

    def set_labels(self, root):
        self.dfs(root, set([]))
        self.total_weight = self.cum_weights[self.root]
        for i, cw in enumerate(self.cum_weights):
            if cw > self.total_weight - cw:
                self.cum_weights[i] = self.total_weight - cw
                self.inverted[i] = True

    def dfs(self, u, visited):
        visited.add(u)
        self.labels[u] = self.nextnode
        self.nextnode += 1
        self.cum_weights[u] = self.weights[u]
        self.num_children[u] = 0
        for v in self.neighbors(u):
            if v not in visited:
                cw, nchildren = self.dfs(v, visited)
                self.cum_weights[u] += cw
                self.num_children[u] += nchildren + 1
        return self.cum_weights[u], self.num_children[u]

    def solve(self):
        tw = self.total_weight
        first_idx = bin_search(self.sorted_nodes, lambda u : self.cum_weights[u] >= (tw+2)//3)
        if first_idx is None:
            return None
        while first_idx < self.size:
            first = self.sorted_nodes[first_idx]
            first_cw = self.cum_weights[first]
            label1 = self.labels[first]
            last_child_label = label1 + self.num_children[first]

            for target_w in (tw - first_cw * 2, first_cw):
                second_idx = bin_search(self.sorted_nodes, lambda u: self.cum_weights[u] >= target_w)
                if second_idx is not None:
                    while (second_idx < self.size
                            and self.cum_weights[self.sorted_nodes[second_idx]] == target_w):
                        second = self.sorted_nodes[second_idx]
                        is_child = label1 < self.labels[second] <= last_child_label
                        is_child = is_child if not self.inverted[first] else not is_child
                        if not is_child and self.labels[second] != label1:
                            return first_cw * 3 - tw
                        second_idx += 1

            first_idx += 1
        if 2*first_cw == tw:
            return first_cw
        return None


def read_input():
    n = int(input())
    weights = r()
    G = WeightedTree(n, weights)
    for _ in range(n-1):
        u, v = r()
        u -= 1
        v -= 1
        G.add_edge(u, v)
    G.initialize()
    return G

def main():
    q = int(input())
    for _ in range(q):
        G = read_input()
        ans = G.solve()
        print(ans if ans is not None else -1)
main()