🥳每日一练-统计二叉树节点的数量-JS简易版

今天分享的内容是如何统计一个排序二叉树节点的数量。

具体的需求是这样的：一个普通的排序二叉树，现在需要在每个节点都加一个 count 属性，表示以该节点为根的子树下面所拥有节点的数量。

如上图，1 节点的 count 为 0，2 节点的 count 为 2，4 节点的 count 为 5 等等

问题很清晰，解决起来不难。具体思路是先计算叶子节点的 count，然后依次往上计算父节点的 count，知道计算完 root 的 count 就结束。

不是排序二叉树，也可以。但看到的算法题目就是这么要求的🐒

准备数据

javascript 复制代码

const data = [9, 4, 5, 3, 2, 7, 8, 95, 33, 22];

const generator = (data) => {
	const insert = (tree, value) => {
		if (!tree) return { value, left: null, right: null};
		if (value < tree.value) {
			tree.left = insert(tree.left, value);
		} else {
			tree.right = insert(tree.right, value);
		}
		return tree;
	};
  
	let tree = null;
	data.forEach((item) => {
		tree = insert(tree, item);
	});
	return tree;
};

const tree = generator(data);

上面代码用generator 函数将 data 数组转成了一个排序二叉树 tree。将 tree 打印出来是这个样子：

以 9 节点为根节点的二叉树。打印的深度为叶子节点下一层，左右子树为空的话，就打印 null。

对二叉树打印感兴趣的话，可以看下面这个代码：

javascript 复制代码

let res = "";

const printTree1 = (data, deeps = [1]) => {
	let space = deeps
		.slice(0, -1)
		.map((item) => {
			return item == 1 ? "|\t\t" : "\t\t";
		})
		.join("");

	space += "|__";
	const content = data ? data.value : null;
	res = res + space + content + "\n";
	
  if (!data) return;
	printTree1(data.left, [...deeps, 1]);
	printTree1(data.right, [...deeps, 0]);

	return res;
};

这个代码不是重点，就不讲解了

整理节点数量

先子节点，后父节点，这是典型的后序遍历啊。那就先复习一遍后序遍历吧

javascript 复制代码

const printTree = (tree) => {
	if (!tree) return null;
	printTree(tree.left);
	printTree(tree.right);
	console.log(tree.value);
};

printTree(tree);
// 2
// 3
// 8
// 7
// 5
// 4
// 22
// 33
// 95
// 9

这是一个简单直观的后序遍历函数。先左子树，后右子树，再根节点。我们可以在函数最后一行的输出做文章，将其修改为可以统计节点数量的功能

javascript 复制代码

const calculateNodeCount = (node) => {
	if (!node) return;
	node.count = (node.left?.count || 0) + (node.right?.count || 0);
	if (node.left) node.count++;
	if (node.right) node.count++;
};

const collectCount = (tree) => {
	if (!tree) return null;
	collectCount(tree.left);
	collectCount(tree.right);
  //calculate Node Count
  calculateNodeCount(tree);
  
	console.log(tree.value + "," + tree.count);
};

collectCount(tree);
// 2,0
// 3,1
// 8,0
// 7,1
// 5,2
// 4,5
// 22,0
// 33,1
// 95,2
// 9,9

输出的结果，前面是节点的值，后面是节点的 count。大家可以和上面打印出来的二叉树树形结构对比，看看 count 值是否正确

逻辑很简单，先计算左子树的 count，然后计算右子树的 count，在统计根的 count。

很简单吧

懂得了计算思路，然后就是直接套用二叉树的遍历模版就可以了。其实关于二叉树的很多问题都是这样，只要你明白了如何遍历，你就能解决问题。甚至是后面将会学习的图的 DFS（深度优先遍历），BFS（广度优先遍历），这些都是基于树的遍历算法上。而 Dijkstra，prim，kruskal，floyd 算法等等，都是基于图的 DFS 和 BFS。所以二叉树的遍历是基础，是掌握复杂数据结构的基石

非递归版本

懂得了递归版本，非递归版本怎么写呢？还是和上面一样，直接套用后序遍历的代码模版即可

javascript 复制代码

const collectCount2 = (tree) => {
	const stack = [];
	let node = tree,
		pre = null;
	while (node || stack.length !== 0) {
		while (node) {
			stack.push(node);
			node = node.left;
		}
		const currentNode = stack.slice(-1)[0];
		if (currentNode.right == null || pre == currentNode.right) {
			//calculate node count
			calculateNodeCount(currentNode);
      console.log(currentNode.value + "," + currentNode.count);
			pre = currentNode;
			stack.pop();
		} else {
			node = currentNode.right;
		}
	}
};

collectCount2(tree);
// 2,0
// 3,1
// 8,0
// 7,1
// 5,2
// 4,5
// 22,0
// 33,1
// 95,2
// 9,9

输出的结果，前面是节点的值，后面是节点的 count。大家可以和上面打印出来的二叉树树形结构对比，看看 count 值是否正确

代码就是简单的二叉树后序遍历的非递归代码，不解释了。想看详细解释的，可以看这篇文章：🥳每日一练-二叉树后序的遍历-JS简易版 - 掘金

总结：

这篇文章分享了如何整理二叉树节点的数量。文中提供了两种解决办法，一种是递归版本，另一种非递归版本。问题不难，理解了解题思路，代码部分就是套用二叉树的遍历模版了

你觉得这篇文章怎么样？我每天都会分享一篇算法小练习，喜欢就点赞+关注吧