华为OD统一考试 - 计算三叉搜索树的高度

题目描述

定义构造三叉搜索树规则如下：

每个节点都存有一个数，当插入一个新的数时，从根节点向下寻找，直到找到一个合适的空节点插入。查找的规则是：

1. 如果数小于节点的数减去500，则将数插入节点的左子树
2. 如果数大于节点的数加上500，则将数插入节点的右子树
3. 否则，将数插入节点的中子树

给你一系列数，请按以上规则，按顺序将数插入树中，构建出一棵三叉搜索树，最后输出树的高度。

输入描述

第一行为一个数 N，表示有 N 个数，1 ≤ N ≤ 10000

第二行为 N 个空格分隔的整数，每个数的范围为[1,10000]

输出描述

输出树的高度（根节点的高度为1）

用例

题目解析

本题应该只是需要模拟出三叉树结构，以及根据指定的逻辑进行插入新节点。

三叉树节点的数据结构定义如下：

{

  val， // 节点值

  height，// 节点所在高度，

  left，// 左子树根节点，

  right，// 右子树根节点，

  mid，// 中子树根节点

}

三叉树的数据结构定义如下：

{

  root，// 根节点

  height，// 树的高度

}

三叉树插入新节点node逻辑：

• 如果三叉树为空，则三叉树根节点root = 新节点node
• 如果三叉树不为空，则从三叉树根节点作为比较节点cur开始比较：

1. 如果 node.val < cur.val - 500，则node应该插入到cur节点的左子树中，若 cur 节点不存在左子树，那么 node 就作为 cur 节点的左子树根节点，且node.height = cur.height + 1若 cur 节点存在左子树，那么 cur = cur.left，继续和 node 比较
2. 如果 node.val > cur.val + 500，则node应该插入到cur节点的右子树中，若 cur 节点不存在右子树，那么 node 就作为 cur 节点的右子树根节点，且node.height = cur.height + 1若 cur 节点存在右子树，那么 cur = cur.right，继续和 node 比较
3. 其他情况，则 node 应该插入到 cur 节点的中子树中，若 cur 节点不存在中子树，那么 node 就作为 cur 节点的中子树根节点，且且node.height = cur.height + 1若 cur 节点存在中子树，那么 cur = cur.mid，继续和 node 比较

在上面比较过程，始终保留最大的node.height作为tree.height。

import Foundation

func ODTest_64() {
    print("第一行为一个数 N，表示有 N 个数，1 ≤ N ≤ 10000")
    let N = Int(readLine() ?? "") ?? 0
    print("第二行为 N 个空格分隔的整数，每个数的范围为[1,10000]")
    let nums = (readLine() ?? "").split(separator: " ").map { Int($0) ?? 0 }
    print("输出树的高度（根节点的高度为1）")

    let tree = Tree()
    for i in 0 ..< N {
        tree.add(nums[i])
    }

    print("\(tree.height)")

    class TreeNode {
        var val = 0 // 节点值
        var height = 0 // 节点所在高度
        var left: TreeNode? // 左子树
        var mid: TreeNode? // 中子树
        var right: TreeNode? // 右子树

        init(_ val: Int) {
            self.val = val
        }
    }

    class Tree {
        var root: TreeNode? // 树的根节点
        var height = 0 // 树的高度

        func add(_ val: Int) {
            var node = TreeNode(val)

            if self.root == nil {
                // 如果树是空的，则当前创建的节点将作为根节点
                node.height = 1 // 根节点的高度为1
                self.root = node
                self.height = 1
            } else {
                // 如果树不是空的，则从根节点开始比较
                var cur = self.root!
                while true {
                    // 假设创建的节点node是当前节点cur的子节点，则node节点高度值=cur节点高度值+1
                    node.height = cur.height + 1
                    // 如果创建的node进入新层，则更新树的高度
                    self.height = max(node.height, self.height)

                    if val < cur.val - 500 {
                        // 如果数小于节点的数减去500，则将数插入cur节点的左子树
                        if cur.left == nil {
                            // 如果cur节点没有左子树，则node作为cur节点的左子树
                            cur.left = node
                            // 停止探索
                            break
                        } else {
                            // 否则继续探索
                            cur = cur.left!
                        }
                    } else if val > cur.val + 500 {
                        // 如果数大于节点的数加上500，则将数插入节点的右子树
                        if cur.right == nil {
                            cur.right = node
                            break
                        } else {
                            cur = cur.right!
                        }
                    } else {
                        // 如果数大于节点的数加上500，则将数插入节点的右子树
                        if cur.mid == nil {
                            cur.mid = node
                            break
                        } else {
                            cur = cur.mid!
                        }
                    }
                }
            }
        }
    }
}