[Algorithm] Universal Value Tree Problem

A unival tree (which stands for "universal value") is a tree where all nodes under it have the same value.

Given the root to a binary tree, count the number of unival subtrees.

For example, the following tree has 5 unival subtrees:

   0
  / \
 1   0
    / \
   1   0
  / \
 1   1

function Node(val) {
  return {
    val,
    left: null,
    right: null
  };
}

const root = Node(0);
root.left = Node(1);
root.right = Node(0);
root.right.left = Node(1);
root.right.right = Node(0);
root.right.left.left = Node(1);
root.right.left.right = Node(1);

function count_unival(root) {
  function helper(root) {
    let total_count = 0;
    let is_unival = true;

    // Base case 1: if current node is null, then return
    if (root == null) {
      return [0, true];
    }

    // Base case 2: if current node is not null, but its children node
    // are null, then count this node as usb-unvial tree
    if (root.left === null && root.right === null) {
      return [1, true];
    }

    // Base case 1 & Base case 2 can keep just one, it should still works

    // Do the Recursion
    let [left_count, is_left_unival] = helper(root.left);
    let [right_count, is_right_unival] = helper(root.right);

    // we need to consider whether the whole tree
    // root + left tree + right tree are unvial
    // the way to do it just compare root with its left and right node
    // whether they are the same if both left tree and right tree are
    // unival tree.
    if (!is_left_unival || !is_right_unival) {
      is_unival = false;
    }

    if (root.left !== null && root.val !== root.left.val) {
      is_unival = false;
    }

    if (root.right !== null && root.val !== root.right.val) {
      is_unival = false;
    }

    // If the whole tree are unival tree, then the final result
    // should + 1
    if (is_unival) {
      return [left_count + right_count + 1, is_unival];
    } else {
      return [left_count + right_count, is_unival];
    }
  }

  const [total_count, is_unival] = helper(root);
  return [total_count, is_unival];
}

const res = count_unival(root);
console.log(
  `Whole tree is${
    res[1] ? "" : "n't"
  } unival tree, total counts for sub-unival tree is ${res[0]}`
); // Whole tree isn't unival tree, total count is 5