Learning Journal

2/10/2024

Count Nodes Equal to Average of Subtree - 2nd attempt

/**
 * Definition for a binary tree node.
 * class TreeNode {
 *     val: number
 *     left: TreeNode | null
 *     right: TreeNode | null
 *     constructor(val?: number, left?: TreeNode | null, right?: TreeNode | null) {
 *         this.val = (val===undefined ? 0 : val)
 *         this.left = (left===undefined ? null : left)
 *         this.right = (right===undefined ? null : right)
 *     }
 * }
 */

function averageOfSubtree(root: TreeNode | null): number {
    const [,, result] = inspect(root);

    return result;
};

function inspect(node: TreeNode):[number, number, number] {
    let sum = node.val, numNodes = 1, numValids = 0;

    if(node.left) {
        const [sL, nL, vL] = inspect(node.left);
        sum += sL;
        numNodes += nL;
        numValids += vL;
    }

    if(node.right) {
        const [sR, nR, vR] = inspect(node.right);
        sum += sR;
        numNodes += nR;
        numValids += vR;
    }

    if(Math.floor(sum/numNodes) === node.val)
        numValids++;
    
    return [sum, numNodes, numValids];
}/**
 * Definition for a binary tree node.
 * class TreeNode {
 *     val: number
 *     left: TreeNode | null
 *     right: TreeNode | null
 *     constructor(val?: number, left?: TreeNode | null, right?: TreeNode | null) {
 *         this.val = (val===undefined ? 0 : val)
 *         this.left = (left===undefined ? null : left)
 *         this.right = (right===undefined ? null : right)
 *     }
 * }
 */

function averageOfSubtree(root: TreeNode | null): number {
    const [,, result] = inspect(root);

    return result;
};

function inspect(node: TreeNode):[number, number, number] {
    let sum = node.val, numNodes = 1, numValids = 0;

    if(node.left) {
        const [sL, nL, vL] = inspect(node.left);
        sum += sL;
        numNodes += nL;
        numValids += vL;
    }

    if(node.right) {
        const [sR, nR, vR] = inspect(node.right);
        sum += sR;
        numNodes += nR;
        numValids += vR;
    }

    if(Math.floor(sum/numNodes) === node.val)
        numValids++;
    
    return [sum, numNodes, numValids];
}

THE PROBLEM

First, this time, instead of letting inspect return only two values, sum and numNodes, I made it return an additional one, numValids indicates the number of valid nodes found in the current subtree. It makes the code significantly shorter (although whether it is clearer is another question) because now, I can shift the conditional check if(Math.floor(sum/numNodes) === node.val) to inspect and in the main function, I only need to call inspect once.

Second, I just recognized something interesting. I don't need to check for null (if(!root)) at the beginning of inspect because I stop right before the function reaches null. The stop is put at two conditions, if(node.left) and if(node.right). So, basically, what I learned from her is that the corner cases depend on where you put the stop, and the stop is usually right before you make another call of that same function.

Revised by ChatGPT

In this approach, I've modified the inspect function to return three values: sum, numNodes, and numValids. The numValids represents the count of valid nodes found within the current subtree. This adjustment significantly shortens the code, as it consolidates the condition if(Math.floor(sum/numNodes) == node.val within the inspect function. Consequently, in the main function, it's necessary to call inspect only once, streamlining the process.

Moreover, I've made an interesting observation: it's not required to check for null using if(!root) at the start of inspect. This is because the function is designed to halt before encountering null, with stops implemented at two conditions: if(node.left) and if(node.right). This strategy illustrates that handling corner cases is contingent upon the placement of these stops, which are strategically positioned just before another recursive call of the same function.

This approach not only simplifies the code but also provides a deeper understanding of managing base cases and recursive function calls in tree traversal algorithms.