C#性能优化-树形结构递归优化

前言

大家好，我是wacky，最近在工作中遇到一个有趣的问题，同事反馈说WPF中有一个树形结构的集合，在加载时会直接报堆栈溢出，一直没时间(懒得)看，导致很久了也没人解决掉。于是，组长就把这个"艰巨"的任务交给了我。作为新人中的"高手"，必然要义不容辞地接受挑战喽，废话不多说，走起。

分析

由于同事此前已经定位到出现问题的代码段，所以到我手中时要省掉不少功夫。打开代码后看了下，原来是这个树形结构使用了典型的递归操作来对每个节点的数据进行更新，在数据量一般时一切正常，但是当数据量达到几万个节点后，这段代码会直接报堆栈溢出的错误。

代码示例如下所示，已简化：

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace Tree
{
    internal class TreeNode
    {
        public int Value { get; set; }
        public List<TreeNode> Children { get; set; }   

        public TreeNode(int value)
        {
            Value = value;
            Children = new List<TreeNode>();
        }
    }
}

// See https://aka.ms/new-console-template for more information
// 创建一个树形结构
using Tree;

internal class Program
{
    static void Main(string[] args)
    {
        TreeNode root = new TreeNode(1);
        TreeNode node2 = new TreeNode(2);
        TreeNode node3 = new TreeNode(3);
        TreeNode node4 = new TreeNode(4);
        TreeNode node5 = new TreeNode(5);
        TreeNode node6 = new TreeNode(6);

        root.Children.Add(node2);
        root.Children.Add(node3);
        node2.Children.Add(node4);
        node2.Children.Add(node5);
        node3.Children.Add(node6);

        PrintTreeNode(root);
        Console.Read();
    }

    static void PrintTreeNode(TreeNode node)
    {
        if (node == null)
        {
            return;
        }
        Console.WriteLine(node.Value);
        foreach (TreeNode child in node.Children)
        {
            PrintTreeNode(child);
        }
    }
}

上述代码我们定义了一个树形结构的类，并加入对应节点，然后使用递归的方式将所有节点输出，在数据量达到前文提到的数量级时就会发生堆栈溢出。

既然是堆栈溢出，那么我们就需要考虑减少堆栈溢出的方式，也就是降低栈的深度。这里我们需要分析下为什么递归会导致堆栈溢出？顺便复习一下部分计算机基础知识点。

在计算机中，函数调用是通过栈(stack)这种数据结构去实现的，每当程序在调用一次函数时，就会进行压栈(push)，每当函数返回后，才会进行出栈(pop)。但是栈的大小本身并不是无限的，加上我们使用C# CLR给的默认分配也不会很大，通常是在1MB左右，这样就会出现函数调用次数过多时，超出栈本身的大小，导致堆栈溢出。

而递归调用，一般都是在到达最后的结束点时，才会一层一层返回每个函数执行的结果。在本次例子中，树形结构存在几万个父子节点，就会导致递归层数过深，函数在栈中无法及时出栈，进而报错。

到这一步时，我们的思路就开始明朗了，既然递归会导致堆栈过深，那我们不妨把递归进行改写，使用其他方式来进行遍历。在通常的解法中，存在两种方式：尾递归优化和迭代。

尾递归优化

什么是尾递归优化？我们先说说什么是尾递归，尾递归是指在一个函数中，所有递归的调用都出现在函数的末尾，也就是递归的那一句在函数执行的最后，或代码路径在最后一句出现，我们就可以称之为尾递归。所以如果我们的递归调用本身不是尾递归的时候，可以通过改写，让它变成尾递归的方式。

为什么尾递归可以进行优化？原因是堆栈需要保存每次调用的返回地址及当时所有的局部变量状态，期间堆栈空间是无法释放的。使用尾递归堆栈可以不用保存上次的函数返回地址/各种状态值，而方法遗留在堆栈上的数据完全可以释放掉，这是尾递归优化的核心思想。

回到我们本次的例子中来，我们的代码已经是尾递归的形式了，但还会导致溢出，那这时我们就需要使用另外一种方法迭代去解决问题了。

迭代

迭代，在本质上就是循环，由于我们已经提到了递归在函数调用的过程中不会对栈进行弹出，那么我们就可以用迭代来模拟入栈出栈的方式来对遍历做优化。我们可以先定义一个栈用来存放所有父子节点，然后对父节点进行压栈，并使用while循环来模拟所有遍历操作，当栈不为空时就一直执行。在循环中我们可以对已经压栈的数据进行弹栈，做完逻辑操作后，再对其子节点进行压栈，一直重复此过程，直到所有节点都弹栈完成。

相关代码如下所示：

// See https://aka.ms/new-console-template for more information
// 创建一个树形结构
using Tree;

internal class Program
{
    static void Main(string[] args)
    {
        TreeNode root = new TreeNode(1);
        TreeNode node2 = new TreeNode(2);
        TreeNode node3 = new TreeNode(3);
        TreeNode node4 = new TreeNode(4);
        TreeNode node5 = new TreeNode(5);
        TreeNode node6 = new TreeNode(6);

        root.Children.Add(node2);
        root.Children.Add(node3);
        node2.Children.Add(node4);
        node2.Children.Add(node5);
        node3.Children.Add(node6);

        IterativeTraversal(root);
        Console.Read();
    }

    static void IterativeTraversal(TreeNode root)
    {
        if (root == null)
        {
            return;
        }
        //定义一个栈，存放所有的树节点
        Stack<TreeNode> stack = new Stack<TreeNode>();
        //把根节点压栈
        stack.Push(root);
        while (stack.Count > 0)
        {
            TreeNode node = stack.Pop();
            Console.WriteLine(node.Value);
            //遍历完父节点后，将子节点压栈
            for (int i = node.Children.Count - 1; i >= 0; i--)
            {
                stack.Push(node.Children[i]);
            }
        }
    }
}

在这种方式中，我们每遍历一层节点，都会对栈进行释放，这样就保证了已经在栈中的层级不会太深，进而解决了堆栈溢出的问题。

总结

探寻好思路后，我和同事做了尝试，将代码改写完成后，遍历几万个节点一切正常，且不会出现卡死之类的其他问题，完美解决！虽然我们本次性能优化的思路并不复杂，代码写起来也相对简单，但背后其实蕴含着比较深刻的计算机原理知识。我们在日常工作中也需要多重视基础知识，包括数据结构和算法，这样才可以在遇到难以解决的问题时游刃有余，诸君共勉！

本文首发于我的公众号【wacky的碎碎念】，喜欢的话可以微信扫码关注哟，我们一起来聊聊技术，谈谈职场和人生~