关于求树中每层的节点总数

关于求树中每层的节点总数

我们用链表来表示树，树也是一种简单的图，同样也可以用二维矩阵来表示。

方案1：基于树的深度优先遍历，进行递归，推荐。

方案2：基于树的深度优先遍历，进行迭代。

方案3：基于树的广度优先遍历，进行迭代。

示例树如下图

代码如下，附带详细注释：

package com.collonn.algorithm.tree;

import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Queue;
import java.util.Set;
import java.util.Stack;

/**
 * 树节点定义，每个节点可以有多个子节点
 */
class TreeLevelNode {
	// 节点值
	private String value;
	// 子节点列表，为了简单，直接用LinkedList
	// 当然，你也可以自定义一个TreeLevelNode链表
	private List<TreeLevelNode> subNodeList = new LinkedList<TreeLevelNode>();

	public TreeLevelNode() {

	}

	public TreeLevelNode(String value) {
		this.value = value;
	}

	public TreeLevelNode(String value, List<TreeLevelNode> subNodeList) {
		this.value = value;
		this.subNodeList = subNodeList;
	}

	public String getValue() {
		return value;
	}

	public void setValue(String value) {
		this.value = value;
	}

	public List<TreeLevelNode> getSubNodeList() {
		return subNodeList;
	}

	public void setSubNodeList(List<TreeLevelNode> subNodeList) {
		this.subNodeList = subNodeList;
	}

	public void addSub(TreeLevelNode node) {
		this.subNodeList.add(node);
	}

	public void addSub(TreeLevelNode[] nodes) {
		for (TreeLevelNode node : nodes) {
			this.subNodeList.add(node);
		}
	}
}

/**
 * <pre>
 * 示例树如下
 * ***********************************************
 * 						a
 * 			b			c			d
 * 		e			f		g		h
 * 								i		j
 * 	***********************************************
 * </pre>
 */
public class TreeLevelCount {
	// 树的根节点
	private TreeLevelNode root;
	// 每层的节点总数
	private Map<Integer, Integer> levelMap = new HashMap<Integer, Integer>();

	private void initTree() {
		// init all nodes
		TreeLevelNode a = new TreeLevelNode("a");
		TreeLevelNode b = new TreeLevelNode("b");
		TreeLevelNode c = new TreeLevelNode("c");
		TreeLevelNode d = new TreeLevelNode("d");
		TreeLevelNode e = new TreeLevelNode("e");
		TreeLevelNode f = new TreeLevelNode("f");
		TreeLevelNode g = new TreeLevelNode("g");
		TreeLevelNode h = new TreeLevelNode("h");
		TreeLevelNode i = new TreeLevelNode("i");
		TreeLevelNode j = new TreeLevelNode("j");

		// set root
		this.root = a;

		// set subNodes
		a.addSub(new TreeLevelNode[] { b, c, d });

		b.addSub(new TreeLevelNode[] { e });
		c.addSub(new TreeLevelNode[] { f, g });
		d.addSub(new TreeLevelNode[] { h });

		h.addSub(new TreeLevelNode[] { i, j });
	}

	// 打印结果
	private void printLevelCount() {
		System.out.printf("------------------- level count\n");
		for (Entry<Integer, Integer> entry : this.levelMap.entrySet()) {
			System.out.printf("level:%d, count:%d\n", entry.getKey(), entry.getValue());
		}
	}

	// 深度优先遍历，递归方式，统计每层的节点数
	// 简单，高效
	private void dfsRecursiveLevelCount(TreeLevelNode node, int level) {
		System.out.printf("visit, node:%s, level:%d\n", node.getValue(), level);

		// 更新计数
		Integer levelCount = this.levelMap.get(level);
		if (levelCount == null) {
			levelCount = 1;
		} else {
			levelCount++;
		}
		this.levelMap.put(level, levelCount);

		// 对每个子节点进行递归调用
		for (TreeLevelNode subNode : node.getSubNodeList()) {
			dfsRecursiveLevelCount(subNode, level + 1);
		}
	}

	// 深度优先遍历，借用栈(非递归)，统计每层的节点数
	// 稍微有点麻烦
	private void dfsStackLevelCount() {
		// 方案1，存储已访问过的节点
		// 方案2，也可以存储某节点已访问过的子节点数，这样可以直接从子节点列表中取出下一个未访问的子节点，这种情况下，用ArrayList存储子节点更高效
		// 方案3，TreeLevelNode中维护一个指针，直接指向同辈中的下一个节点
		Set<String> nodeValueSet = new HashSet<String>();
		// 用来保存当前路径的栈
		Stack<TreeLevelNode> stack = new Stack<TreeLevelNode>();
		// 初始加入root节点
		stack.add(this.root);

		// 根节点的level永远是1，levelCount永远是1
		this.levelMap.put(1, 1);
		// 将root节点加入已访问
		nodeValueSet.add(this.root.getValue());
		System.out.printf("visit, node:%s, level:%d\n", this.root.getValue(), 1);

		out: while (!stack.isEmpty()) {
			// 访问栈顶元素，但不弹出
			TreeLevelNode topNode = stack.peek();

			// 当其没有子节点时，弹出栈顶元素
			if (topNode.getSubNodeList().size() == 0) {
				stack.pop();
				continue;
			}

			// 寻找一个该节点的未访问的子节点
			for (TreeLevelNode nod : topNode.getSubNodeList()) {
				if (!nodeValueSet.contains(nod.getValue())) {
					// 如果子节点没有被访问过，则加入栈顶
					stack.push(nod);
					// 将该节点已访问
					nodeValueSet.add(nod.getValue());
					System.out.printf("visit, node:%s, level:%d\n", nod.getValue(), stack.size());

					// 注意，栈的当前大小，就是栈顶元素在树中的层次
					// 更新level的统计总数
					Integer count = this.levelMap.get(stack.size());
					if (count == null) {
						count = 1;
					} else {
						count++;
					}
					this.levelMap.put(stack.size(), count);

					// 断续深度遍历，相当于调用了递归方法
					continue out;
				}
			}

			// 如果该节点的所有孩子节点都访问过，则弹出栈顶元素
			stack.pop();
		}
	}

	// 广度优先遍历，借用队列，统计每层的节点数
	// 稍微有点麻烦
	private void bfsLevelCount() {
		// 存储元素的层级数
		Map<String, Integer> nodeLevelMap = new HashMap<String, Integer>();
		// 队列
		Queue<TreeLevelNode> queue = new LinkedList<TreeLevelNode>();

		// 初始加入root节点
		queue.offer(this.root);
		// 根节点的level永远是1，levelCount永远是1
		this.levelMap.put(1, 1);
		nodeLevelMap.put(this.root.getValue(), 1);
		System.out.printf("visit, node:%s, level:%d\n", this.root.getValue(), 1);

		while (!queue.isEmpty()) {
			// 访问队列头，并删除
			TreeLevelNode head = queue.poll();

			// 将子节点依次加入到队列尾
			for (TreeLevelNode nd : head.getSubNodeList()) {
				queue.offer(nd);

				// 子节点所在的层次为：父节点层级加1
				int levelCurrent = nodeLevelMap.get(head.getValue()) + 1;
				// 记录节点的层级
				nodeLevelMap.put(nd.getValue(), levelCurrent);
				System.out.printf("visit, node:%s, level:%d\n", nd.getValue(), levelCurrent);

				// 更新level的统计总数
				Integer count = this.levelMap.get(levelCurrent);
				if (count == null) {
					count = 1;
				} else {
					count++;
				}
				this.levelMap.put(levelCurrent, count);
			}
		}
	}

	public static void main(String[] args) {
		TreeLevelCount tlc = new TreeLevelCount();
		tlc.initTree();

		System.out.printf("------------------- 深度优先遍历，递归求解 -------------------\n");
		tlc.levelMap.clear();
		tlc.dfsRecursiveLevelCount(tlc.root, 1);
		tlc.printLevelCount();

		System.out.printf("------------------- 深度优先遍历，迭代求解 -------------------\n");
		tlc.levelMap.clear();
		tlc.dfsStackLevelCount();
		tlc.printLevelCount();

		System.out.printf("------------------- 广度优先遍历 -------------------\n");
		tlc.levelMap.clear();
		tlc.bfsLevelCount();
		tlc.printLevelCount();
	}

}

代码执行结果如下：

------------------- 深度优先遍历，递归求解 -------------------
visit, node:a, level:1
visit, node:b, level:2
visit, node:e, level:3
visit, node:c, level:2
visit, node:f, level:3
visit, node:g, level:3
visit, node:d, level:2
visit, node:h, level:3
visit, node:i, level:4
visit, node:j, level:4
------------------- level count
level:1, count:1
level:2, count:3
level:3, count:4
level:4, count:2
------------------- 深度优先遍历，迭代求解 -------------------
visit, node:a, level:1
visit, node:b, level:2
visit, node:e, level:3
visit, node:c, level:2
visit, node:f, level:3
visit, node:g, level:3
visit, node:d, level:2
visit, node:h, level:3
visit, node:i, level:4
visit, node:j, level:4
------------------- level count
level:1, count:1
level:2, count:3
level:3, count:4
level:4, count:2
------------------- 广度优先遍历 -------------------
visit, node:a, level:1
visit, node:b, level:2
visit, node:c, level:2
visit, node:d, level:2
visit, node:e, level:3
visit, node:f, level:3
visit, node:g, level:3
visit, node:h, level:3
visit, node:i, level:4
visit, node:j, level:4
------------------- level count
level:1, count:1
level:2, count:3
level:3, count:4
level:4, count:2

原创博文，转载请注明出处。