LeetCode-1382.平衡二叉搜索树代码优化解决TLE问题咨询

核心性能瓶颈 #

• 高度计算冗余：原有height函数每次递归遍历整棵子树算高度，updateHeight插入后全树更新高度，单次插入开销O(n)，整体复杂度O(n²)，是超时的核心原因。
• 不平衡节点查找冗余：每次插入后全树遍历找不平衡节点，实际上只有插入路径上的节点平衡因子会变化，无需全树扫描。
• 旋转实现冗余：LLRotation中额外malloc节点副本，旋转只需修改指针即可，多余的内存操作严重拖慢性能还会造成内存泄漏。
• 基础错误：createNode未初始化height字段，malloc时类型写错（sizeof(struct TreeNode)应为sizeof(struct TreeNodeC)），旋转逻辑仅支持LL场景，缺少另外三种旋转实现。

优化后代码 #

#include <stdlib.h>
#include <stdbool.h>

// 题目给定的原TreeNode定义，此处仅为编译通过补充
struct TreeNode {
    int val;
    struct TreeNode *left;
    struct TreeNode *right;
};

struct TreeNodeC {
    int val;
    int height;
    struct TreeNodeC *left;
    struct TreeNodeC *right;
};

struct TreeNodeC *avl = NULL;

int maximum(int a, int b) {
    return a > b ? a : b;
}

// 直接取节点保存的高度，无需递归计算
int height(struct TreeNodeC *root) {
    return root == NULL ? 0 : root->height;
}

// 计算平衡因子，保留正负判断倾斜方向
int get_balance_factor(struct TreeNodeC *root) {
    return root == NULL ? 0 : height(root->left) - height(root->right);
}

struct TreeNodeC *createNode(int val) {
    struct TreeNodeC *newNode = (struct TreeNodeC *)malloc(sizeof(struct TreeNodeC));
    newNode->left = NULL;
    newNode->right = NULL;
    newNode->val = val;
    // 新建节点高度初始为1
    newNode->height = 1;
    return newNode;
}

// LL右旋
struct TreeNodeC *rightRotate(struct TreeNodeC *y) {
    struct TreeNodeC *x = y->left;
    struct TreeNodeC *T2 = x->right;

    // 执行旋转
    x->right = y;
    y->left = T2;

    // 仅更新涉及节点的高度，无需全树更新
    y->height = maximum(height(y->left), height(y->right)) + 1;
    x->height = maximum(height(x->left), height(x->right)) + 1;

    return x;
}

// RR左旋
struct TreeNodeC *leftRotate(struct TreeNodeC *x) {
    struct TreeNodeC *y = x->right;
    struct TreeNodeC *T2 = y->left;

    // 执行旋转
    y->left = x;
    x->right = T2;

    // 更新高度
    x->height = maximum(height(x->left), height(x->right)) + 1;
    y->height = maximum(height(y->left), height(y->right)) + 1;

    return y;
}

// 递归插入，回溯时自动更新高度、检查平衡
struct TreeNodeC *insert(struct TreeNodeC *node, int val) {
    // 1. 正常BST插入
    if (node == NULL)
        return createNode(val);

    if (val < node->val)
        node->left = insert(node->left, val);
    else if (val > node->val)
        node->right = insert(node->right, val);
    else // 重复值不插入
        return node;

    // 2. 回溯更新当前节点高度
    node->height = 1 + maximum(height(node->left), height(node->right));

    // 3. 检查当前节点是否失衡
    int balance = get_balance_factor(node);

    // 4. 失衡则执行对应旋转，共四种情况
    // LL
    if (balance > 1 && val < node->left->val)
        return rightRotate(node);
    // RR
    if (balance < -1 && val > node->right->val)
        return leftRotate(node);
    // LR
    if (balance > 1 && val > node->left->val) {
        node->left = leftRotate(node->left);
        return rightRotate(node);
    }
    // RL
    if (balance < -1 && val < node->right->val) {
        node->right = rightRotate(node->right);
        return leftRotate(node);
    }

    // 未失衡直接返回当前节点
    return node;
}

void avlTreeInsertion(int val) {
    avl = insert(avl, val);
}

void inOrderTraversal(struct TreeNode *root) {
    if (root != NULL) {
        inOrderTraversal(root->left);
        avlTreeInsertion(root->val);
        inOrderTraversal(root->right);
    }
}

struct TreeNodeC *balanceBST(struct TreeNode *root) {
    avl = NULL; // 每次调用先重置全局变量，避免用例间干扰
    inOrderTraversal(root);
    return avl;
}

优化效果说明 #

所有插入过程中的高度更新、平衡检查、旋转操作都仅在插入路径上执行，单次插入开销为O(log n)，整体时间复杂度从原有O(n²)降到O(n log n)，完全适配LeetCode大数据量用例的时间要求，同时修正了原有代码的内存泄漏、初始化错误、旋转逻辑不全的问题。

内容的提问来源于stack exchange，提问作者T.Rakesh