Problem: Minimum Score After Removals on a Tree

There is an undirected connected tree with `n` nodes labeled from `0` to `n -
1` and `n - 1` edges.

You are given a 0-indexed integer array `nums` of length `n` where `nums[i]`
represents the value of the `ith` node. You are also given a 2D integer array
`edges` of length `n - 1` where `edges[i] = [ai, bi]` indicates that there is
an edge between nodes `ai` and `bi` in the tree.

Remove two distinct edges of the tree to form three connected components. For
a pair of removed edges, the following steps are defined:

1. Get the XOR of all the values of the nodes for each of the three components respectively.
2. The difference between the largest XOR value and the smallest XOR value is the score of the pair.

* For example, say the three components have the node values: `[4,5,7]`, `[1,9]`, and `[3,3,3]`. The three XOR values are `4 ^ 5 ^ 7 = 6`, `1 ^ 9 = 8`, and `3 ^ 3 ^ 3 = 3`. The largest XOR value is `8` and the smallest XOR value is `3`. The score is then `8 - 3 = 5`.

Return the minimum score of any possible pair of edge removals on the given
tree.


Python solution:
import math
import itertools
import bisect
import collections
import string
import heapq
import functools
import sortedcontainers
from typing import List, Dict, Tuple, Iterator, Set

class Solution:
  def minimumScore(self, nums: List[int], edges: List[List[int]]) -> int:
    n = len(nums)
    xors = functools.reduce(lambda x, y: x ^ y, nums)
    subXors = nums[:]
    tree = [[] for _ in range(n)]
    children = [{i} for i in range(n)]

    for u, v in edges:
      tree[u].append(v)
      tree[v].append(u)

    def dfs(u: int, parent: int) -> Tuple[int, Set[int]]:
      for v in tree[u]:
        if v == parent:
          continue
        vXor, vChildren = dfs(v, u)
        subXors[u] ^= vXor
        children[u] |= vChildren
      return subXors[u], children[u]

    dfs(0, -1)

    ans = math.inf
    for i in range(len(edges)):
      a, b = edges[i]
      if b in children[a]:
        a, b = b, a
      for j in range(i):
        c, d = edges[j]
        if d in children[c]:
          c, d = d, c

        if c in children[a] and a != c:
          cands = [subXors[c], subXors[a] ^ subXors[c], xors ^ subXors[a]]
        elif a in children[c] and a != c:
          cands = [subXors[a], subXors[c] ^ subXors[a], xors ^ subXors[c]]
        else:
          cands = [subXors[a], subXors[c], xors ^ subXors[a] ^ subXors[c]]
        ans = min(ans, max(cands) - min(cands))

    return ans


Java solution:
class Solution {
  public int minimumScore(int[] nums, int[][] edges) {
    final int n = nums.length;
    final int xors = getXors(nums);
    int[] subXors = nums.clone();
    List<Integer>[] tree = new List[n];
    Set<Integer>[] children = new Set[n];

    for (int i = 0; i < n; ++i)
      tree[i] = new ArrayList<>();

    for (int i = 0; i < n; ++i)
      children[i] = new HashSet<>(Arrays.asList(i));

    for (int[] edge : edges) {
      final int u = edge[0];
      final int v = edge[1];
      tree[u].add(v);
      tree[v].add(u);
    }

    dfs(tree, 0, -1, subXors, children);

    int ans = Integer.MAX_VALUE;

    for (int i = 0; i < edges.length; ++i) {
      int a = edges[i][0];
      int b = edges[i][1];
      if (children[a].contains(b)) {
        final int temp = a;
        a = b;
        b = a;
      }
      for (int j = 0; j < i; ++j) {
        int c = edges[j][0];
        int d = edges[j][1];
        if (children[c].contains(d)) {
          final int temp = c;
          c = d;
          d = temp;
        }
        int[] cands;
        if (a != c && children[a].contains(c))
          cands = new int[] {subXors[c], subXors[a] ^ subXors[c], xors ^ subXors[a]};
        else if (a != c && children[c].contains(a))
          cands = new int[] {subXors[a], subXors[c] ^ subXors[a], xors ^ subXors[c]};
        else
          cands = new int[] {subXors[a], subXors[c], xors ^ subXors[a] ^ subXors[c]};
        ans = Math.min(ans, Arrays.stream(cands).max().getAsInt() -
                                Arrays.stream(cands).min().getAsInt());
      }
    }

    return ans;
  }

  private Pair<Integer, Set<Integer>> dfs(List<Integer>[] tree, int u, int prev, int[] subXors,
                                          Set<Integer>[] children) {
    for (final int v : tree[u]) {
      if (v == prev)
        continue;
      final Pair<Integer, Set<Integer>> pair = dfs(tree, v, u, subXors, children);
      final int vXor = pair.getKey();
      final Set<Integer> vChildren = pair.getValue();
      subXors[u] ^= vXor;
      for (final int child : vChildren)
        children[u].add(child);
    }
    return new Pair<>(subXors[u], children[u]);
  }

  private int getXors(int[] nums) {
    int xors = 0;
    for (final int num : nums)
      xors ^= num;
    return xors;
  }
}


C++ solution:
class Solution {
 public:
  int minimumScore(vector<int>& nums, vector<vector<int>>& edges) {
    const int n = nums.size();
    const int xors = reduce(nums.begin(), nums.end(), 0, bit_xor());
    vector<int> subXors(nums);
    vector<vector<int>> tree(n);
    vector<unordered_set<int>> children(n);

    for (int i = 0; i < n; ++i)
      children[i].insert(i);

    for (const vector<int>& edge : edges) {
      const int u = edge[0];
      const int v = edge[1];
      tree[u].push_back(v);
      tree[v].push_back(u);
    }

    dfs(tree, 0, -1, subXors, children);

    int ans = INT_MAX;

    for (int i = 0; i < edges.size(); ++i) {
      int a = edges[i][0];
      int b = edges[i][1];
      if (children[a].contains(b))
        swap(a, b);
      for (int j = 0; j < i; ++j) {
        int c = edges[j][0];
        int d = edges[j][1];
        if (children[c].contains(d))
          swap(c, d);
        vector<int> cands;
        if (a != c && children[a].contains(c))
          cands = {subXors[c], subXors[a] ^ subXors[c], xors ^ subXors[a]};
        else if (a != c && children[c].contains(a))
          cands = {subXors[a], subXors[c] ^ subXors[a], xors ^ subXors[c]};
        else
          cands = {subXors[a], subXors[c], xors ^ subXors[a] ^ subXors[c]};
        ans = min(ans, ranges::max(cands) - ranges::min(cands));
      }
    }

    return ans;
  }

 private:
  pair<int, unordered_set<int>> dfs(const vector<vector<int>>& tree, int u,
                                    int prev, vector<int>& subXors,
                                    vector<unordered_set<int>>& children) {
    for (const int v : tree[u]) {
      if (v == prev)
        continue;
      const auto& [vXor, vChildren] = dfs(tree, v, u, subXors, children);
      subXors[u] ^= vXor;
      children[u].insert(vChildren.begin(), vChildren.end());
    }
    return {subXors[u], children[u]};
  }
};