7-1 Is It An AVL Tree (PAT ADSAA)

#include 
#include 
#include 

struct node {
  int key;
  node* leftNode;
  node* rightNode;
  int height;
};

node* buildTree(std::vector vec) {
  node* root;
  root = new node;
  root->key = vec[0];
  auto it = vec.begin() + 1;
  if (it == vec.end()) {
    root->leftNode = nullptr;
    root->rightNode = nullptr;
    root->height = 1;
    return root;
  }
  while (it != vec.end()) {
    if (*it > vec[0]) {
      break;
    }
    ++it;
  }
  if (it == vec.begin() + 1) {
    root->leftNode = nullptr;
    std::vector right(it, vec.end());
    root->rightNode = buildTree(right);
    root->height = 1 + root->rightNode->height;
  } else if (it == vec.end()) {
    root->rightNode = nullptr;
    std::vector left(vec.begin() + 1, vec.end());
    root->leftNode = buildTree(left);
    root->height = 1 + root->leftNode->height;
  } else {
    std::vector left(vec.begin() + 1, it);
    std::vector right(it, vec.end());
    root->leftNode = buildTree(left);
    root->rightNode = buildTree(right);
    root->height =
        1 + std::max(root->leftNode->height, root->rightNode->height);
  }
  return root;
}

bool judgeAVL(node* root) {
  if (root->leftNode && root->rightNode) {
    int gap = root->leftNode->height - root->rightNode->height;
    if (gap >= -1 && gap <= 1 && judgeAVL(root->leftNode) &&
        judgeAVL(root->rightNode)) {
      return true;
    }
    return false;
  }
  if (!root->leftNode && !root->rightNode) {
    return true;
  }
  if (root->leftNode && !root->rightNode) {
    if (root->leftNode->height <= 1 && judgeAVL(root->leftNode)) {
      return true;
    }
    return false;
  }
  if (!root->leftNode && root->rightNode) {
    if (root->rightNode->height <= 1 && judgeAVL(root->rightNode)) {
      return true;
    }
    return false;
  }
}

int main() {
  int K, N;
  std::cin >> K;
  for (int i = 0; i < K; ++i) {
    std::cin >> N;
    std::vector vec(N);
    for (int j = 0; j < N; ++j) {
      std::cin >> vec[j];
    }
    node* root;
    root = buildTree(vec);
    if (judgeAVL(root)) {
      std::cout << "Yes" << std::endl;
    } else {
      std::cout << "No" << std::endl;
    }
  }
  return 0;
}

题目如下：

In computer science, an AVL tree (Georgy Adelson-Velsky and Evgenii Landis' tree, named after the inventors) is a self-balancing binary search tree. It was the first such data structure to be invented. In an AVL tree, the heights of the two child subtrees of any node differ by at most one. (Quoted from wikipedia)

For each given binary search tree, you are supposed to tell if it is an AVL tree.

Input Specification:

Each input file contains several test cases. The first line gives a positive integer K (≤10) which is the total number of cases. For each case, the first line gives a positive integer N (≤30), the total number of nodes in the binary search tree. The second line gives the preorder traversal sequence of the tree with all the keys being distinct. All the numbers in a line are separated by a space.

Output Specification:

For each test case, print in a line "Yes" if the given tree is an AVL tree, or "No" if not.

Sample Input:

3
7
50 40 36 48 46 62 77
8
50 40 36 48 46 62 77 88
6
50 40 36 48 46 62

Sample Output:

Yes
No
No