#include "rb_tree.h"
#include <library/cpp/testing/unittest/registar.h>
#include <util/random/fast.h>
#include <util/random/easy.h>
#include <util/random/shuffle.h>
class TRedBlackTreeTest: public TTestBase {
struct TCmp {
template <class T>
static inline bool Compare(const T& l, const T& r) {
return l.N < r.N;
}
template <class T>
static inline bool Compare(const T& l, int r) {
return l.N < r;
}
template <class T>
static inline bool Compare(int l, const T& r) {
return l < r.N;
}
};
class TNode: public TRbTreeItem<TNode, TCmp> {
public:
inline TNode(int n) noexcept
: N(n)
{
}
int N;
};
using TTree = TRbTree<TNode, TCmp>;
UNIT_TEST_SUITE(TRedBlackTreeTest);
UNIT_TEST(TestEmpty)
UNIT_TEST(TestInsert)
UNIT_TEST(TestErase)
UNIT_TEST(TestFind)
UNIT_TEST(TestStress)
UNIT_TEST(TestGettingIndexWithDifferentValues)
UNIT_TEST(TestCheckChildrenAfterErase)
UNIT_TEST(TestGettingIndexWithDifferentValuesAfterErase)
UNIT_TEST(TestGettingIndexWithEqualValues)
UNIT_TEST(TestLessCountOnEmptyTree)
UNIT_TEST_SUITE_END();
private:
inline void TestStress() {
TVector<TSimpleSharedPtr<TNode>> nodes;
for (int i = 0; i < 1000; ++i) {
nodes.push_back(new TNode(i));
}
TTree tree;
TReallyFastRng32 rnd(Random());
for (size_t i = 0; i < 1000000; ++i) {
tree.Insert(nodes[rnd.Uniform(nodes.size())].Get());
}
for (TTree::TConstIterator it = tree.Begin(); it != tree.End();) {
const int v1 = it->N;
if (++it == tree.End()) {
break;
}
const int v2 = it->N;
UNIT_ASSERT(v1 < v2);
}
}
inline void TestGettingIndexWithDifferentValues() {
TVector<TSimpleSharedPtr<TNode>> nodes;
size_t N = 1000;
for (size_t i = 0; i < N; ++i) {
nodes.push_back(new TNode(int(i)));
}
TTree tree;
Shuffle(nodes.begin(), nodes.end());
for (size_t i = 0; i < N; ++i) {
tree.Insert(nodes[i].Get());
}
for (size_t i = 0; i < N; ++i) {
UNIT_ASSERT_EQUAL(tree.LessCount(i), i);
UNIT_ASSERT_EQUAL(tree.NotGreaterCount(i), i + 1);
UNIT_ASSERT_EQUAL(tree.GreaterCount(i), N - i - 1);
UNIT_ASSERT_EQUAL(tree.NotLessCount(i), N - i);
auto nodePtr = tree.Find(i);
UNIT_ASSERT_EQUAL(tree.GetIndex(nodePtr), i);
UNIT_ASSERT_EQUAL(tree.GetIndex(nodes[i].Get()), static_cast<size_t>(nodes[i]->N));
}
}
inline void TestCheckChildrenAfterErase() {
TVector<TSimpleSharedPtr<TNode>> nodes;
size_t N = 1000;
for (size_t i = 0; i < N; ++i) {
nodes.push_back(new TNode(int(i)));
}
TTree tree;
Shuffle(nodes.begin(), nodes.end());
for (size_t i = 0; i < N; ++i) {
tree.Insert(nodes[i].Get());
}
auto checker = [](const TTree& tree) {
for (auto node = tree.Begin(); node != tree.End(); ++node) {
size_t childrens = 1;
if (node->Left_) {
childrens += node->Left_->Children_;
}
if (node->Right_) {
childrens += node->Right_->Children_;
}
UNIT_ASSERT_VALUES_EQUAL(childrens, node->Children_);
}
};
for (auto node : nodes) {
tree.Erase(node.Get());
checker(tree);
}
}
inline void TestGettingIndexWithDifferentValuesAfterErase() {
TVector<TSimpleSharedPtr<TNode>> nodes;
size_t N = 1000;
for (size_t i = 0; i < N; ++i) {
nodes.push_back(new TNode(int(i)));
}
TTree tree;
Shuffle(nodes.begin(), nodes.end());
for (size_t i = 0; i < N; ++i) {
tree.Insert(nodes[i].Get());
}
{
size_t index = 0;
for (auto node = tree.Begin(); node != tree.End(); ++node, ++index) {
UNIT_ASSERT_VALUES_EQUAL(tree.GetIndex(&*node), index);
UNIT_ASSERT_VALUES_EQUAL(tree.ByIndex(index)->N, node->N);
UNIT_ASSERT_VALUES_EQUAL(node->N, index);
}
}
for (size_t i = 1; i < N; i += 2) {
auto* node = tree.Find(i);
UNIT_ASSERT_VALUES_EQUAL(node->N, i);
tree.Erase(node);
}
{
size_t index = 0;
for (auto node = tree.Begin(); node != tree.End(); ++node, ++index) {
UNIT_ASSERT_VALUES_EQUAL(tree.GetIndex(&*node), index);
UNIT_ASSERT_VALUES_EQUAL(tree.ByIndex(index)->N, node->N);
UNIT_ASSERT_VALUES_EQUAL(node->N, 2 * index);
}
}
}
inline void TestGettingIndexWithEqualValues() {
TVector<TSimpleSharedPtr<TNode>> nodes;
size_t N = 1000;
for (size_t i = 0; i < N; ++i) {
nodes.push_back(new TNode(0));
}
TTree tree;
for (size_t i = 0; i < N; ++i) {
tree.Insert(nodes[i].Get());
}
for (size_t i = 0; i < N; ++i) {
UNIT_ASSERT_EQUAL(tree.LessCount(nodes[i]->N), 0);
UNIT_ASSERT_EQUAL(tree.NotGreaterCount(nodes[i]->N), N);
UNIT_ASSERT_EQUAL(tree.GreaterCount(nodes[i]->N), 0);
UNIT_ASSERT_EQUAL(tree.NotLessCount(nodes[i]->N), N);
UNIT_ASSERT_EQUAL(tree.LessCount(*nodes[i].Get()), 0);
UNIT_ASSERT_EQUAL(tree.NotGreaterCount(*nodes[i].Get()), N);
UNIT_ASSERT_EQUAL(tree.GreaterCount(*nodes[i].Get()), 0);
UNIT_ASSERT_EQUAL(tree.NotLessCount(*nodes[i].Get()), N);
}
}
inline void TestFind() {
TTree tree;
{
TNode n1(1);
TNode n2(2);
TNode n3(3);
tree.Insert(n1);
tree.Insert(n2);
tree.Insert(n3);
UNIT_ASSERT_EQUAL(tree.Find(1)->N, 1);
UNIT_ASSERT_EQUAL(tree.Find(2)->N, 2);
UNIT_ASSERT_EQUAL(tree.Find(3)->N, 3);
UNIT_ASSERT(!tree.Find(0));
UNIT_ASSERT(!tree.Find(4));
UNIT_ASSERT(!tree.Find(1234567));
}
UNIT_ASSERT(tree.Empty());
}
inline void TestEmpty() {
TTree tree;
UNIT_ASSERT(tree.Empty());
UNIT_ASSERT_EQUAL(tree.Begin(), tree.End());
}
inline void TestInsert() {
TTree tree;
{
TNode n1(1);
TNode n2(2);
TNode n3(3);
tree.Insert(n1);
tree.Insert(n2);
tree.Insert(n3);
TTree::TConstIterator it = tree.Begin();
UNIT_ASSERT_EQUAL((it++)->N, 1);
UNIT_ASSERT_EQUAL((it++)->N, 2);
UNIT_ASSERT_EQUAL((it++)->N, 3);
UNIT_ASSERT_EQUAL(it, tree.End());
}
UNIT_ASSERT(tree.Empty());
}
inline void TestErase() {
TTree tree;
{
TNode n1(1);
TNode n2(2);
TNode n3(3);
tree.Insert(n1);
tree.Insert(n2);
tree.Insert(n3);
TTree::TIterator it = tree.Begin();
tree.Erase(it++);
UNIT_ASSERT_EQUAL(it, tree.Begin());
UNIT_ASSERT_EQUAL(it->N, 2);
tree.Erase(it++);
UNIT_ASSERT_EQUAL(it, tree.Begin());
UNIT_ASSERT_EQUAL(it->N, 3);
tree.Erase(it++);
UNIT_ASSERT_EQUAL(it, tree.Begin());
UNIT_ASSERT_EQUAL(it, tree.End());
}
UNIT_ASSERT(tree.Empty());
}
inline void TestLessCountOnEmptyTree() {
TTree tree;
UNIT_ASSERT_VALUES_EQUAL(0, tree.LessCount(TNode(1)));
}
};
UNIT_TEST_SUITE_REGISTRATION(TRedBlackTreeTest);