【PAT甲级】1066. Root of AVL Tree (25)

#include <stdio.h>
#include <math.h>
#include <stdlib.h>

typedef struct Node {
    Node *left;
    Node *right;
    int key;
} Node;

void insert(Node* &r, int k);
int height(Node *r);
int max(int a, int b);
void balance(Node* &n);
void rotate_left(Node* &n);
void rotate_right(Node* &n);

int main(int argc, char *argv[]) {
    int m;
    Node *tree = NULL;
    scanf("%d", &m);
    for (int i = 0; i < m; i++) {
        int k;
        scanf("%d", &k);
        insert(tree, k);
    }
    printf("%d\n", tree->key);
    return 0;
}

void insert(Node* &r, int k) {
    if (r == NULL) {
        Node *n = (Node *)malloc(sizeof(Node));
        n->key = k;
        n->left = n->right = NULL;
        r = n;
    } else if (k < r->key) {
        insert(r->left, k);
        if (abs(height(r->left)-height(r->right)) > 1)
            balance(r);
    } else {
        insert(r->right, k);
        if (abs(height(r->left)-height(r->right)) > 1)
            balance(r);
    }
    return;
}
int height(Node *r) {
    if (r == NULL) return 0;
    return max(height(r->left), height(r->right)) + 1;
}
int max(int a, int b) {
    return a > b ? a : b;
}
void balance(Node* &n) {
    if (height(n->left) - height(n->right) == 2) {
        if (height(n->left->left) > height(n->left->right)) {
            rotate_right(n);
        } else {
            rotate_left(n->left);
            rotate_right(n);
        }
    } else if (height(n->left) - height(n->right) == -2) {
        if (height(n->right->right) > height(n->right->left)) {
            rotate_left(n);
        } else {
            rotate_right(n->right);
            rotate_left(n);
        }
    }
    return;
}
void rotate_right(Node* &n) {
    Node *a = n->left;
    n->left = n->left->right;
    a->right = n;
    n = a;
    return;
}
void rotate_left(Node* &n) {
    Node *a = n->right;
    n->right = n->right->left;
    a->left = n;
    n = a;
    return;
}