线索二叉树

发布时间 2023-06-24 10:12:42作者: hacker_dvd

题目链接:普通二叉树(简化版)

#include <bits/stdc++.h>

template <typename T>
class ThreadedBinaryTreeNode {
public:
  T data;
  int count;
  ThreadedBinaryTreeNode* left;
  ThreadedBinaryTreeNode* right;
  bool leftThread;
  bool rightThread;

  ThreadedBinaryTreeNode(T value)
    : data(value), count(1), left(nullptr), right(nullptr), leftThread(true), rightThread(true) {}
};

template <typename T>
class ThreadedBinaryTree {
public:
  ThreadedBinaryTreeNode<T>* root;

  ThreadedBinaryTree() : root(nullptr) {}
  // 中序遍历
  void inorderTraversal();
  // 查找前驱
  ThreadedBinaryTreeNode<T>* findPredecessor(T value);
  // 查找后继
  ThreadedBinaryTreeNode<T>* findSuccessor(T value);
  // 查找值为value的节点的排名
  int findRank(T value);
  // 查找排名为rank的节点的值
  T findByRank(int rank);
  // 添加一个值为value的节点
  void add(T value);
  void add(ThreadedBinaryTreeNode<T>* node, T value);
  // 查找最左边的节点
  ThreadedBinaryTreeNode<T>* findLeftMostNode(ThreadedBinaryTreeNode<T>* node);
  // 查找最右边的节点
  ThreadedBinaryTreeNode<T>* findRightMostNode(ThreadedBinaryTreeNode<T>* node);
  // 查找值为value的节点
  ThreadedBinaryTreeNode<T>* findNode(T value);
  // 计算以node为根的子树的节点数
  int countNodes(ThreadedBinaryTreeNode<T>* node);
};

template <typename T>
void ThreadedBinaryTree<T>::inorderTraversal() {
  ThreadedBinaryTreeNode<T>* current = findLeftMostNode(root);
  while (current != nullptr) {
    for (int i = 0; i < current->count; ++i) {
      std::cout << current->data << " ";
    }
    if (current->rightThread) {
      current = current->right;
    } else {
      current = findLeftMostNode(current->right);
    }
  }
  std::cout << std::endl;
}

template <typename T>
ThreadedBinaryTreeNode<T>* ThreadedBinaryTree<T>::findPredecessor(T value) {
  ThreadedBinaryTreeNode<T>* current = root;
  ThreadedBinaryTreeNode<T>* predecessor = nullptr;
  while (current != nullptr) {
    if (value < current->data) {
      if (current->leftThread) {
        break;
      }
      current = current->left;
    } else if (value > current->data) {
      predecessor = current;
      if (current->rightThread) {
        break;
      }
      current = current->right;
    } else {
      if (!current->leftThread) {
        predecessor = findRightMostNode(current->left);
      }
      break;
    }
  }
  return predecessor;
}

template <typename T>

ThreadedBinaryTreeNode<T>* ThreadedBinaryTree<T>::findSuccessor(T value) {
  ThreadedBinaryTreeNode<T>* current = root;
  ThreadedBinaryTreeNode<T>* successor = nullptr;
  while (current != nullptr) {
    if (value < current->data) {
      successor = current;
      if (current->leftThread) {
        break;
      }
      current = current->left;
    } else if (value > current->data) {
      if (current->rightThread) {
        break;
      }
      current = current->right;
    } else {
      if (!current->rightThread) {
        successor = findLeftMostNode(current->right);
      }
      break;
    }
  }
  return successor;
}

template <typename T>
int ThreadedBinaryTree<T>::findRank(T value) {
  auto tnode = findLeftMostNode(root);
  if (tnode && tnode->data > value) {
    return 1;
  }
  tnode = findRightMostNode(root);
  if (tnode && tnode->data < value) {
    return countNodes(root) + 1;
  }
  int rank = 0;
  ThreadedBinaryTreeNode<T>* current = root;
  while (current != nullptr) {
    if (value < current->data) {
      if (current->leftThread) {
        rank++;
        break;
      }
      current = current->left;
    } else {
      if (!current->leftThread) {
        rank += countNodes(current->left) + current->count;
      } else {
        rank += current->count;
      }
      if (value == current->data) {
        rank = rank - current->count + 1;
        break;
      }
      if (current->rightThread) {
        rank++;
        break;
      }
      current = current->right;
    }
  }
  return rank;
}

template <typename T>
T ThreadedBinaryTree<T>::findByRank(int rank) {
  int currentRank = 0;
  ThreadedBinaryTreeNode<T>* current = root;
  while (current != nullptr) {
    if (!current->leftThread) {
      int leftCount = countNodes(current->left);
      if (currentRank + leftCount + current->count > rank) {
        if (currentRank + leftCount < rank) {
          return current->data;
        }
        current = current->left;
      } else if (currentRank + leftCount + current->count < rank) {
        currentRank += leftCount + current->count;
        current = current->right;
      } else {
        return current->data;
      }
    } else {
      if (currentRank + current->count >= rank) {
        return current->data;
      }
      currentRank += current->count;
      current = current->right;
    }
  }
  return T(); // 返回默认值,如果找不到对应排名的数
}

template <typename T>
void ThreadedBinaryTree<T>::add(T value) {
  if (root == nullptr) {
    root = new ThreadedBinaryTreeNode<T>(value);
  } else {
    add(root, value);
  }
}

template <typename T>
void ThreadedBinaryTree<T>::add(ThreadedBinaryTreeNode<T>* node, T value) {
  if (value < node->data) {
    if (node->leftThread || node->left == nullptr) {
      ThreadedBinaryTreeNode<T>* newNode = new ThreadedBinaryTreeNode<T>(value);
      newNode->left = node->left;
      newNode->right = node;
      node->left = newNode;
      node->leftThread = false;
    } else {
      add(node->left, value);
    }
  } else if (value > node->data) {
    if (node->rightThread || node->right == nullptr) {
      ThreadedBinaryTreeNode<T>* newNode = new ThreadedBinaryTreeNode<T>(value);
      newNode->left = node;
      newNode->right = node->right;
      node->right = newNode;
      node->rightThread = false;
    } else {
      add(node->right, value);
    }
  } else {
    node->count++;
  }
}

template <typename T>
ThreadedBinaryTreeNode<T>* ThreadedBinaryTree<T>::findLeftMostNode(ThreadedBinaryTreeNode<T>* node) {
  if (node == nullptr) {
    return nullptr;
  }
  while (!node->leftThread && node->left != nullptr) {
    node = node->left;
  }
  return node;
}

template <typename T>
ThreadedBinaryTreeNode<T>* ThreadedBinaryTree<T>::findRightMostNode(ThreadedBinaryTreeNode<T>* node) {
  if (node == nullptr) {
    return nullptr;
  }
  while (!node->rightThread && node->right != nullptr) {
    node = node->right;
  }
  return node;
}

template <typename T>
ThreadedBinaryTreeNode<T>* ThreadedBinaryTree<T>::findNode(T value) {
  ThreadedBinaryTreeNode<T>* current = root;
  while (current != nullptr) {
    if (value < current->data) {
      if (current->leftThread) {
        break;
      }
      current = current->left;
    } else if (value > current->data) {
      if (current->rightThread) {
        break;
      }
      current = current->right;
    } else {
      return current;
    }
  }
  return nullptr;
}

template <typename T>
int ThreadedBinaryTree<T>::countNodes(ThreadedBinaryTreeNode<T>* node) {
  if (node == nullptr) {
    return 0;
  }
  int count = node->count;
  if (!node->leftThread) {
    count += countNodes(node->left);
  }
  if (!node->rightThread) {
    count += countNodes(node->right);
  }
  return count;
}


using namespace std;
const int INF = 2147483647;

int main() {
  int q;
  cin >> q;
  ThreadedBinaryTree<int> tree;
  while (q--) {
    int op, x;
    cin >> op >> x;
    switch (op) {
    case 1: {
      printf("%d\n", tree.findRank(x));
      break;
    }
    case 2: {
      printf("%d\n", tree.findByRank(x));
      break;
    }
    case 3: {
      auto pre = tree.findPredecessor(x);
      printf("%d\n", pre == nullptr ? -INF : pre->data);
      break;
    }
    case 4: {
      auto suc = tree.findSuccessor(x);
      printf("%d\n", suc == nullptr ? INF : suc->data);
      break;
    }
    case 5: {
      tree.add(x);
      break;
    }
    }
  }
}