444 lines
No EOL
11 KiB
C++
444 lines
No EOL
11 KiB
C++
#pragma once
|
|
|
|
#include "ContainersCommon.hpp"
|
|
|
|
namespace tp {
|
|
|
|
template <typename NumericType>
|
|
struct AvlNumericKey {
|
|
public:
|
|
AvlNumericKey() = default;
|
|
AvlNumericKey(NumericType val) :
|
|
val(val) {}
|
|
|
|
inline bool descentRight(const AvlNumericKey& in) const { return in.val > val; }
|
|
inline bool exactNode(const AvlNumericKey& in) const { return in.val == val; }
|
|
|
|
inline const AvlNumericKey& keyInRightSubtree(const AvlNumericKey& in) const { return in; }
|
|
inline const AvlNumericKey& keyInLeftSubtree(const AvlNumericKey& in) const { return in; }
|
|
|
|
template <typename NodeType>
|
|
inline void updateNodeCache(const NodeType*) {}
|
|
|
|
template <typename NodeType>
|
|
inline const AvlNumericKey& getFindKey(const NodeType*) const {
|
|
return *this;
|
|
}
|
|
|
|
public:
|
|
NumericType val;
|
|
};
|
|
|
|
template <typename Key, typename Data, class Allocator = DefaultAllocator>
|
|
class AvlTree {
|
|
typedef SelectValueOrReference<Key> KeyArg;
|
|
typedef SelectValueOrReference<Data> DataArg;
|
|
|
|
public:
|
|
class Node {
|
|
friend AvlTree;
|
|
|
|
private:
|
|
Node(KeyArg aKey, DataArg aData) :
|
|
key(aKey),
|
|
data(aData) {}
|
|
|
|
public:
|
|
Data data;
|
|
Key key;
|
|
|
|
Node* mLeft = nullptr;
|
|
Node* mRight = nullptr;
|
|
Node* mParent = nullptr;
|
|
ualni mHeight = 0;
|
|
};
|
|
|
|
public:
|
|
AvlTree() = default;
|
|
~AvlTree() { removeAll(); }
|
|
|
|
[[nodiscard]] ualni size() const { return mSize; }
|
|
|
|
Node* head() const { return this->mRoot; }
|
|
|
|
void insert(KeyArg key, DataArg data) {
|
|
mRoot = insertUtil(mRoot, key, data);
|
|
mRoot->mParent = nullptr;
|
|
}
|
|
|
|
void remove(KeyArg key) {
|
|
mRoot = removeUtil(mRoot, key);
|
|
if (mRoot) mRoot->mParent = nullptr;
|
|
}
|
|
|
|
Node* maxNode(Node* head) const {
|
|
if (!head) return nullptr;
|
|
while (head->mRight != nullptr) {
|
|
head = head->mRight;
|
|
}
|
|
return head;
|
|
}
|
|
|
|
Node* minNode(Node* head) const {
|
|
if (!head) return nullptr;
|
|
while (head->mLeft != nullptr) {
|
|
head = head->mLeft;
|
|
}
|
|
return head;
|
|
}
|
|
|
|
Node* find(KeyArg key) const {
|
|
Node* iter = mRoot;
|
|
while (true) {
|
|
if (!iter) return nullptr;
|
|
if (iter->key.exactNode(key)) return iter;
|
|
if (iter->key.descentRight(key)) {
|
|
key = iter->key.keyInRightSubtree(key);
|
|
iter = iter->mRight;
|
|
} else {
|
|
key = iter->key.keyInLeftSubtree(key);
|
|
iter = iter->mLeft;
|
|
}
|
|
}
|
|
}
|
|
|
|
Node* findLessOrEq(KeyArg key) const {
|
|
Node* iter = mRoot;
|
|
while (true) {
|
|
if (!iter) return nullptr;
|
|
if (iter->key.exactNode(key)) return iter;
|
|
if (iter->key.descentRight(key)) {
|
|
if (iter->mRight) {
|
|
key = iter->key.keyInRightSubtree(key);
|
|
iter = iter->mRight;
|
|
} else {
|
|
return iter;
|
|
}
|
|
} else {
|
|
if (iter->mLeft) {
|
|
key = iter->key.keyInLeftSubtree(key);
|
|
iter = iter->mLeft;
|
|
} else {
|
|
return iter;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// checks invariants of AVL tree
|
|
// returns first invalid node
|
|
const Node* findInvalidNode(const Node* head) const {
|
|
if (head == nullptr) return nullptr;
|
|
|
|
if (head->mLeft) {
|
|
// TODO: incomplete test
|
|
if (head->key.descentRight(head->mLeft->key.getFindKey(head))) {
|
|
return head;
|
|
}
|
|
if (head->mLeft->mParent != head) {
|
|
return head;
|
|
}
|
|
if (!head->mRight && head->mLeft->mHeight != head->mHeight - 1) {
|
|
return head;
|
|
}
|
|
}
|
|
|
|
if (head->mRight) {
|
|
if (!head->key.descentRight(head->mRight->key.getFindKey(head))) {
|
|
return head;
|
|
}
|
|
if (head->mRight->mParent != head) {
|
|
return head;
|
|
}
|
|
if (!head->mLeft && head->mRight->mHeight != head->mHeight - 1) {
|
|
return head;
|
|
}
|
|
}
|
|
|
|
if (head->mLeft && head->mRight) {
|
|
if (max(head->mLeft->mHeight, head->mRight->mHeight) != head->mHeight - 1) {
|
|
return head;
|
|
}
|
|
}
|
|
|
|
int balance = getNodeHeight(head->mRight) - getNodeHeight(head->mLeft);
|
|
|
|
if (balance > 1 || balance < -1) {
|
|
return head;
|
|
}
|
|
|
|
const Node* ret = findInvalidNode(head->mRight);
|
|
|
|
if (ret) {
|
|
return ret;
|
|
}
|
|
|
|
return findInvalidNode(head->mLeft);
|
|
}
|
|
|
|
bool isValid() { return findInvalidNode(head()) == nullptr; }
|
|
|
|
template <typename tFunctor>
|
|
void traverse(Node* node, bool after, tFunctor functor) {
|
|
if (!after) functor(node);
|
|
if (node->mLeft) traverse(node->mLeft, after, functor);
|
|
if (node->mRight) traverse(node->mRight, after, functor);
|
|
if (after) functor(node);
|
|
}
|
|
|
|
void removeAll() {
|
|
if (!mRoot) return;
|
|
removeUtil(mRoot);
|
|
mRoot = nullptr;
|
|
mSize = 0;
|
|
}
|
|
|
|
void removeUtil(Node* node) {
|
|
if (node->mLeft) removeUtil(node->mLeft);
|
|
if (node->mRight) removeUtil(node->mRight);
|
|
deleteNode(node);
|
|
}
|
|
|
|
public:
|
|
template <class tArchiver>
|
|
void archiveWrite(tArchiver& file) const {
|
|
FAIL("not implemented")
|
|
}
|
|
|
|
template <class tArchiver>
|
|
void archiveRead(tArchiver&) {
|
|
FAIL("not implemented")
|
|
}
|
|
|
|
private:
|
|
inline void deleteNode(Node* node) {
|
|
node->~Node();
|
|
mAlloc.deallocate(node);
|
|
mSize--;
|
|
}
|
|
|
|
inline Node* newNode(KeyArg key, DataArg data) { return new (mAlloc.allocate(sizeof(Node))) Node(key, data); }
|
|
|
|
// swaps pointers of two nodes to preserve data location in the memory instead of swapping data directly
|
|
inline void injectNodeInstead(Node* target, Node* from) {
|
|
// 'target' always has two nodes
|
|
// 'from' has only one child on the right on no child at all
|
|
// 'from' can be right child of the 'target'
|
|
// 'target' can or can not have parent
|
|
// 'target' can be left or right child of parent
|
|
|
|
// from is not a child of target
|
|
Node* targetParent = target->mParent;
|
|
Node* targetLeft = target->mLeft;
|
|
Node* targetRight = target->mRight;
|
|
|
|
bool special = target->mRight == from;
|
|
|
|
// update parent
|
|
if (targetParent) {
|
|
if (targetParent->mRight == target) targetParent->mRight = from;
|
|
else targetParent->mLeft = from;
|
|
}
|
|
|
|
// clone 'from' to 'target'
|
|
target->mParent = special ? from : from->mParent;
|
|
target->mLeft = nullptr;
|
|
target->mRight = from->mRight;
|
|
if (!special) target->mParent->mLeft = target;
|
|
|
|
// if (target->mLeft) target->mLeft->mParent = target;
|
|
if (target->mRight) target->mRight->mParent = target;
|
|
|
|
// clone 'target' to 'from'
|
|
from->mParent = targetParent;
|
|
from->mLeft = targetLeft;
|
|
from->mRight = special ? target : targetRight;
|
|
from->mLeft->mParent = from;
|
|
from->mRight->mParent = from;
|
|
|
|
std::swap(from->mHeight, target->mHeight);
|
|
}
|
|
|
|
inline alni getNodeHeight(const Node* node) const { return node ? node->mHeight : -1; }
|
|
|
|
// returns new head
|
|
inline Node* rotateLeft(Node* pivot) {
|
|
DEBUG_ASSERT(pivot);
|
|
|
|
Node* const head = pivot;
|
|
Node* const right = pivot->mRight;
|
|
Node* const right_left = right->mLeft;
|
|
Node* const parent = pivot->mParent;
|
|
|
|
// parents
|
|
if (right_left) right_left->mParent = head;
|
|
head->mParent = right;
|
|
right->mParent = parent;
|
|
|
|
// children
|
|
head->mRight = right_left;
|
|
right->mLeft = head;
|
|
|
|
// heights
|
|
head->mHeight = 1 + max(getNodeHeight(head->mLeft), getNodeHeight(head->mRight));
|
|
right->mHeight = 1 + max(getNodeHeight(right->mLeft), getNodeHeight(right->mRight));
|
|
|
|
// cache
|
|
head->key.updateNodeCache(head);
|
|
right->key.updateNodeCache(right);
|
|
|
|
return right;
|
|
}
|
|
|
|
inline Node* rotateRight(Node* pivot) {
|
|
DEBUG_ASSERT(pivot);
|
|
|
|
Node* const head = pivot;
|
|
Node* const left = pivot->mLeft;
|
|
Node* const left_right = left->mRight;
|
|
Node* const parent = pivot->mParent;
|
|
|
|
// parents
|
|
if (left_right) left_right->mParent = head;
|
|
head->mParent = left;
|
|
left->mParent = parent;
|
|
|
|
// children
|
|
head->mLeft = left_right;
|
|
left->mRight = head;
|
|
|
|
// heights
|
|
head->mHeight = 1 + max(getNodeHeight(head->mLeft), getNodeHeight(head->mRight));
|
|
left->mHeight = 1 + max(getNodeHeight(left->mLeft), getNodeHeight(left->mRight));
|
|
|
|
// cache
|
|
head->key.updateNodeCache(head);
|
|
left->key.updateNodeCache(left);
|
|
|
|
return left;
|
|
}
|
|
|
|
// recursively returns valid isLeft or isRight child or root
|
|
Node* insertUtil(Node* head, KeyArg key, DataArg data) {
|
|
|
|
Node* insertedNode;
|
|
|
|
if (head == nullptr) {
|
|
mSize++;
|
|
Node* out = newNode(key, data);
|
|
out->key.updateNodeCache(out);
|
|
return out;
|
|
} else if (head->key.exactNode(key)) {
|
|
return head;
|
|
} else if (head->key.descentRight(key)) {
|
|
insertedNode = insertUtil(head->mRight, head->key.keyInRightSubtree(key), data);
|
|
head->mRight = insertedNode;
|
|
insertedNode->mParent = head;
|
|
} else {
|
|
insertedNode = insertUtil(head->mLeft, head->key.keyInLeftSubtree(key), data);
|
|
head->mLeft = insertedNode;
|
|
insertedNode->mParent = head;
|
|
}
|
|
|
|
// update height
|
|
head->mHeight = 1 + max(getNodeHeight(head->mRight), getNodeHeight(head->mLeft));
|
|
|
|
alni balance = alni(getNodeHeight(head->mRight) - getNodeHeight(head->mLeft));
|
|
|
|
if (balance > 1) {
|
|
if (head->mRight->key.descentRight(head->key.keyInRightSubtree(key))) {
|
|
return rotateLeft(head);
|
|
} else {
|
|
head->mRight = rotateRight(head->mRight);
|
|
return rotateLeft(head);
|
|
}
|
|
} else if (balance < -1) {
|
|
if (!head->mLeft->key.descentRight(head->key.keyInLeftSubtree(key))) {
|
|
return rotateRight(head);
|
|
} else {
|
|
head->mLeft = rotateLeft(head->mLeft);
|
|
return rotateRight(head);
|
|
}
|
|
}
|
|
|
|
head->key.updateNodeCache(head);
|
|
|
|
return head;
|
|
}
|
|
|
|
Node* removeUtil(Node* head, KeyArg key) {
|
|
if (head == nullptr) return head;
|
|
|
|
if (head->key.exactNode(key)) {
|
|
if (head->mRight && head->mLeft) {
|
|
Node* min = minNode(head->mRight);
|
|
injectNodeInstead(head, min);
|
|
std::swap(head, min);
|
|
// auto const& newKey = min->key.getFindKey(head->mRight);
|
|
head->mRight = removeUtil(head->mRight, key);
|
|
} else if (head->mRight) {
|
|
|
|
if (head->mParent) {
|
|
if (head->mParent->mLeft == head) head->mParent->mLeft = head->mRight;
|
|
else head->mParent->mRight = head->mRight;
|
|
}
|
|
|
|
head->mRight->mParent = head->mParent;
|
|
|
|
auto delNode = head;
|
|
head = head->mRight;
|
|
deleteNode(delNode);
|
|
} else if (head->mLeft) {
|
|
|
|
if (head->mParent) {
|
|
if (head->mParent->mLeft == head) head->mParent->mLeft = head->mLeft;
|
|
else head->mParent->mRight = head->mLeft;
|
|
}
|
|
|
|
head->mLeft->mParent = head->mParent;
|
|
|
|
auto delNode = head;
|
|
head = head->mLeft;
|
|
deleteNode(delNode);
|
|
} else {
|
|
deleteNode(head);
|
|
head = nullptr;
|
|
}
|
|
} else if (head->key.descentRight(key)) {
|
|
head->mRight = removeUtil(head->mRight, head->key.keyInRightSubtree(key));
|
|
} else {
|
|
head->mLeft = removeUtil(head->mLeft, head->key.keyInLeftSubtree(key));
|
|
}
|
|
|
|
if (head == nullptr) return head;
|
|
|
|
head->mHeight = 1 + max(getNodeHeight(head->mRight), getNodeHeight(head->mLeft));
|
|
alni balance = getNodeHeight(head->mRight) - getNodeHeight(head->mLeft);
|
|
|
|
if (balance < -1) {
|
|
if (getNodeHeight(head->mLeft->mLeft) >= getNodeHeight(head->mLeft->mRight)) {
|
|
return rotateRight(head);
|
|
} else {
|
|
head->mLeft = rotateLeft(head->mLeft);
|
|
return rotateRight(head);
|
|
}
|
|
} else if (balance > 1) {
|
|
if (getNodeHeight(head->mRight->mRight) >= getNodeHeight(head->mRight->mLeft)) {
|
|
return rotateLeft(head);
|
|
} else {
|
|
head->mRight = rotateRight(head->mRight);
|
|
return rotateLeft(head);
|
|
}
|
|
}
|
|
|
|
head->key.updateNodeCache(head);
|
|
|
|
return head;
|
|
}
|
|
|
|
private:
|
|
Node* mRoot = nullptr;
|
|
ualni mSize = 0;
|
|
Allocator mAlloc;
|
|
};
|
|
} |