tdlight/tdutils/td/utils/DecTree.h
Arseny Smirnov 093651fb22 Update tdutils from other project
GitOrigin-RevId: 83ec94032ccecef812b01963ac0506655a57e2af
2018-08-13 20:15:09 +03:00

178 lines
5.0 KiB
C++

//
// Copyright Aliaksei Levin (levlam@telegram.org), Arseny Smirnov (arseny30@gmail.com) 2014-2018
//
// Distributed under the Boost Software License, Version 1.0. (See accompanying
// file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
//
#pragma once
#include <memory>
#include <utility>
#include "int_types.h"
#include "Random.h"
namespace td {
template <typename KeyType, typename ValueType, typename Compare = std::less<KeyType>>
class DecTree {
private:
struct Node {
std::unique_ptr<Node> left_;
std::unique_ptr<Node> right_;
size_t size_;
KeyType key_;
ValueType value_;
uint32 y_;
void relax() {
size_ = 1;
if (left_ != nullptr) {
size_ += left_->size_;
}
if (right_ != nullptr) {
size_ += right_->size_;
}
}
Node(KeyType key, ValueType value, uint32 y) : key_(std::move(key)), value_(std::move(value)), y_(y) {
size_ = 1;
}
};
std::unique_ptr<Node> root_;
std::unique_ptr<Node> create_node(KeyType key, ValueType value, uint32 y) {
return std::make_unique<Node>(std::move(key), std::move(value), y);
}
std::unique_ptr<Node> insert_node(std::unique_ptr<Node> Tree, KeyType key, ValueType value, uint32 y) {
if (Tree == nullptr) {
return create_node(std::move(key), std::move(value), y);
}
if (Tree->y_ < y) {
auto P = split_node(std::move(Tree), key);
auto T = create_node(std::move(key), std::move(value), y);
T->left_ = std::move(P.first);
T->right_ = std::move(P.second);
T->relax();
return std::move(T);
}
if (Compare()(key, Tree->key_)) {
Tree->left_ = insert_node(std::move(Tree->left_), std::move(key), std::move(value), y);
} else if (Compare()(Tree->key_, key)) {
Tree->right_ = insert_node(std::move(Tree->right_), std::move(key), std::move(value), y);
} else {
// ?? assert
}
Tree->relax();
return std::move(Tree);
}
std::unique_ptr<Node> remove_node(std::unique_ptr<Node> Tree, KeyType &key) {
if (Tree == nullptr) {
// ?? assert
return nullptr;
}
if (Compare()(key, Tree->key_)) {
Tree->left_ = remove_node(std::move(Tree->left_), key);
} else if (Compare()(Tree->key_, key)) {
Tree->right_ = remove_node(std::move(Tree->right_), key);
} else {
Tree = merge_node(std::move(Tree->left_), std::move(Tree->right_));
}
if (Tree != nullptr) {
Tree->relax();
}
return std::move(Tree);
}
ValueType *get_node(std::unique_ptr<Node> &Tree, KeyType &key) {
if (Tree == nullptr) {
return nullptr;
}
if (Compare()(key, Tree->key_)) {
return get_node(Tree->left_, key);
} else if (Compare()(Tree->key_, key)) {
return get_node(Tree->right_, key);
} else {
return &Tree->value_;
}
}
ValueType *get_node_by_idx(std::unique_ptr<Node> &Tree, size_t idx) {
CHECK(Tree != nullptr);
auto s = (Tree->left_ != nullptr) ? Tree->left_->size_ : 0;
if (idx < s) {
return get_node_by_idx(Tree->left_, idx);
} else if (idx == s) {
return &Tree->value_;
} else {
return get_node_by_idx(Tree->right_, idx - s - 1);
}
}
std::pair<std::unique_ptr<Node>, std::unique_ptr<Node>> split_node(std::unique_ptr<Node> Tree, KeyType &key) {
if (Tree == nullptr) {
return std::pair<std::unique_ptr<Node>, std::unique_ptr<Node>>(nullptr, nullptr);
}
if (Compare()(key, Tree->key_)) {
auto P = split_node(std::move(Tree->left_), key);
Tree->left_ = std::move(P.second);
Tree->relax();
P.second = std::move(Tree);
return std::move(P);
} else {
auto P = split_node(std::move(Tree->right_), key);
Tree->right_ = std::move(P.first);
Tree->relax();
P.first = std::move(Tree);
return std::move(P);
}
}
std::unique_ptr<Node> merge_node(std::unique_ptr<Node> left, std::unique_ptr<Node> right) {
if (left == nullptr) {
return std::move(right);
}
if (right == nullptr) {
return std::move(left);
}
if (left->y_ < right->y_) {
right->left_ = merge_node(std::move(left), std::move(right->left_));
right->relax();
return std::move(right);
} else {
left->right_ = merge_node(std::move(left->right_), std::move(right));
left->relax();
return std::move(left);
}
}
public:
DecTree() {
}
size_t size() const {
if (root_ == nullptr) {
return 0;
} else {
return root_->size_;
}
}
void insert(KeyType key, ValueType value) {
root_ = insert_node(std::move(root_), std::move(key), std::move(value), td::Random::fast_uint32());
}
void remove(KeyType &key) {
root_ = remove_node(std::move(root_), key);
}
ValueType *get(KeyType &key) {
return get_node(root_, key);
}
ValueType *get_random() {
if (size() == 0) {
return nullptr;
} else {
return get_node_by_idx(root_, td::Random::fast_uint32() % size());
}
}
bool exists(KeyType &key) {
return get_node(root_, key) != nullptr;
}
};
} // namespace td