LeetCode/java/cousins-in-binary-tree-ii.java

class Solution {
  private int get(TreeNode node) {
    if (node == null) {
      return 0;
    }
    return node.val;
  }

  private void handleNode(Queue<TreeNode> q, int siblings, TreeNode node) {
    if (node == null) {
      return;
    }

    node.val = siblings;
    q.offer(node);
  }

  private int handleLevel(Queue<TreeNode> q, int previous) {
    int sum = 0;

    for (int i = q.size(); i > 0; --i) {
      var node = q.poll();
      node.val = previous - node.val;

      var siblings = get(node.left) + get(node.right);
      handleNode(q, siblings, node.left);
      handleNode(q, siblings, node.right);

      sum += siblings;
    }

    return sum;
  }

  public TreeNode replaceValueInTree(TreeNode root) {
    if (root == null) {
      return root;
    }

    Queue<TreeNode> q = new ArrayDeque<>();

    int sumAbove = root.val;
    q.offer(root);

    while (!q.isEmpty()) {
      sumAbove = handleLevel(q, sumAbove);
    }

    return root;
  }
}