【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;
}