Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- using System.CodeDom.Compiler;
- using System.Collections.Generic;
- using System.Collections;
- using System.ComponentModel;
- using System.Diagnostics.CodeAnalysis;
- using System.Globalization;
- using System.IO;
- using System.Linq;
- using System.Reflection;
- using System.Runtime.Serialization;
- using System.Text.RegularExpressions;
- using System.Text;
- using System;
- class Solution {
- class Node {
- public int Index;
- public int Coin;
- public List<int> Child;
- public int Parent;
- public long SumToParent;
- public HashSet<long> Set;
- public Node(int index, int coin) {
- Index = index;
- Coin = coin;
- Child = new List<int>();
- Parent = -1;
- SumToParent = 0;
- }
- public override string ToString() {
- return $"I = {Index}, C = {Coin}, S = {SumToParent}, P = {Parent}, Ch = " + string.Join(", ", Child.Select(i => i.ToString()));
- }
- }
- class Pair {
- public long Sum;
- public HashSet<long> Set;
- public Pair(long sum, HashSet<long> set) {
- Sum = sum;
- Set = set;
- }
- }
- static void PrintTree(IEnumerable<Node> nodes) {
- foreach (var node in nodes) {
- Console.WriteLine(node);
- }
- }
- static void PrintTree(IEnumerable<int> indices, List<Node> nodes) {
- foreach (var node in indices) {
- Console.WriteLine(nodes[node]);
- }
- }
- // Complete the balancedForest function below.
- static long balancedForest(int[] c, int[][] edges) {
- if (c.Length == 1)
- return -1;
- //Console.WriteLine("*******************************");
- List<Node> nodes = new List<Node>(c.Length);
- long totalSum = 0;
- for (var i = 0; i < c.Length; i++) {
- totalSum += (long)c[i];
- nodes.Add(new Node(i, c[i]));
- }
- foreach (var edge in edges) {
- var x = edge[0] - 1;
- var y = edge[1] - 1;
- nodes[x].Child.Add(y);
- nodes[y].Child.Add(x);
- }
- //PrintTree(nodes);
- HashSet<long> set = new HashSet<long>();
- List<int> levels = new List<int>(c.Length);
- Queue<int> queue = new Queue<int>();
- queue.Enqueue(0);
- levels.Add(0);
- while (queue.Count > 0) {
- var current = nodes[queue.Dequeue()];
- foreach (var child in current.Child) {
- if (child == current.Parent) {
- continue;
- }
- nodes[child].Parent = current.Index;
- levels.Add(child);
- queue.Enqueue(child);
- }
- }
- //PrintTree(levels, nodes);
- levels.Reverse();
- long ans = long.MaxValue;
- foreach (var node in levels) {
- var current = nodes[node];
- current.Set = new HashSet<long>();
- var childSums = new List<Pair>(current.Child.Count - 1);
- foreach (var child in current.Child) {
- if (child == current.Parent) {
- continue;
- }
- childSums.Add(new Pair(nodes[child].SumToParent, nodes[child].Set));
- current.Set.UnionWith(nodes[child].Set);
- nodes[child].Set = null;
- current.SumToParent += nodes[child].SumToParent;
- }
- current.SumToParent += current.Coin;
- current.Set.Add(current.SumToParent);
- childSums.Sort((p1, p2) => p1.Sum.CompareTo(p2.Sum));
- childSums.Reverse();
- //Console.WriteLine($"Current = {current}");
- //Console.WriteLine($"ChildSums : {String.Join(", ", childSums.Select(p => p.Sum))}");
- //Console.WriteLine($"HashSets : {String.Join("; ",
- //childSums.Select(p => String.Join(", ", p.Set)))}"
- //);
- foreach (var pair in childSums) {
- var childSum = pair.Sum;
- var childSet = pair.Set;
- bool wasAdd = set.Add(childSum);
- if (!wasAdd && 3 * childSum >= totalSum) {
- ans = Math.Min(3 * childSum - totalSum, ans);
- Console.WriteLine($"ans = {ans} <-- 1");
- }
- var rem = totalSum - childSum;
- //Console.WriteLine($"childSum = {childSum}, rem = {rem}");
- if (childSet.Contains(rem) && rem * 2 >= childSum) {
- ans = Math.Min(rem * 2 - childSum, ans);
- //Console.WriteLine($"ans = {ans} <-- 2");
- }
- if (rem % 2 == 0 && rem / 2 >= childSum && set.Contains(rem / 2) && !(wasAdd && rem /2 == childSum)) {
- ans = Math.Min(rem / 2 - childSum, ans);
- //Console.WriteLine($"ans = {ans} <-- 3");
- }
- if (rem < childSum && rem * 2 >= childSum && childSet.Contains(childSum - rem)) {
- ans = Math.Min(rem * 2 - childSum, ans);
- //Console.WriteLine($"ans = {ans} <-- 4");
- }
- if (rem > childSum && rem < 2 * childSum && set.Contains(rem - childSum) && !childSet.Contains(rem - childSum)) {
- ans = Math.Min(childSum - (rem - childSum), ans);
- //Console.WriteLine($"ans = {ans} <-- 5");
- }
- }
- }
- if (ans == long.MaxValue) {
- ans = -1;
- }
- return ans;
- }
- static void Main(string[] args) {
- TextWriter textWriter = new StreamWriter(@System.Environment.GetEnvironmentVariable("OUTPUT_PATH"), true);
- int q = Convert.ToInt32(Console.ReadLine());
- for (int qItr = 0; qItr < q; qItr++) {
- int n = Convert.ToInt32(Console.ReadLine());
- int[] c = Array.ConvertAll(Console.ReadLine().Split(' '), cTemp => Convert.ToInt32(cTemp))
- ;
- int[][] edges = new int[n - 1][];
- for (int i = 0; i < n - 1; i++) {
- edges[i] = Array.ConvertAll(Console.ReadLine().Split(' '), edgesTemp => Convert.ToInt32(edgesTemp));
- }
- long result = balancedForest(c, edges);
- textWriter.WriteLine(result);
- }
- textWriter.Flush();
- textWriter.Close();
- }
- }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement