diff --git a/include/util/trbtree.h b/include/util/trbtree.h index d666cc11d1b42a60c07a5bea2926362bcbaa6941..1e34839cbedba7a22e9dd59f90fd03b2e41a33eb 100644 --- a/include/util/trbtree.h +++ b/include/util/trbtree.h @@ -38,8 +38,8 @@ SRBTreeNode *tRBTreeDropByKey(SRBTree *pTree, void *pKey); SRBTreeNode *tRBTreeGet(SRBTree *pTree, void *pKey); // SRBTreeIter ============================================= -#define tRBTreeIterCreate(tree, descend) \ - (SRBTreeIter) { .des = (descend), .pTree = (tree), .pNode = (descend) ? (tree)->maxNode : (tree)->minNode } +#define tRBTreeIterCreate(tree, ascend) \ + (SRBTreeIter) { .asc = (ascend), .pTree = (tree), .pNode = (ascend) ? (tree)->minNode : (tree)->maxNode } SRBTreeNode *tRBTreeIterNext(SRBTreeIter *pIter); @@ -60,7 +60,7 @@ struct SRBTree { }; struct SRBTreeIter { - int8_t des; + int8_t asc; SRBTree *pTree; SRBTreeNode *pNode; }; diff --git a/source/util/src/trbtree.c b/source/util/src/trbtree.c index 14b841dc50b5421ec9e673efebd6ae7023aec9e6..7364c129bbf85da9ab8022883aed14b4bc8b5fbd 100644 --- a/source/util/src/trbtree.c +++ b/source/util/src/trbtree.c @@ -60,6 +60,55 @@ static void tRBTreeRotateRight(SRBTree *pTree, SRBTreeNode *pNode) { pNode->parent = left; } +static SRBTreeNode *tRBTreeSuccessor(SRBTreeNode *pNode) { + if (pNode->right) { + pNode = pNode->right; + while (pNode->left) { + pNode = pNode->left; + } + } else { + while (true) { + if (pNode->parent) { + if (pNode == pNode->parent->left) { + pNode = pNode->parent; + break; + } else { + pNode = pNode->parent; + } + } else { + pNode = NULL; + break; + } + } + } + + return pNode; +} + +static SRBTreeNode *tRBTreePredecessor(SRBTreeNode *pNode) { + if (pNode->left) { + pNode = pNode->left; + while (pNode->right) { + pNode = pNode->right; + } + } else { + while (true) { + if (pNode->parent) { + if (pNode == pNode->parent->right) { + pNode = pNode->parent; + break; + } else { + pNode = pNode->parent; + } + } else { + pNode = NULL; + break; + } + } + } + return NULL; +} + SRBTreeNode *tRBTreePut(SRBTree *pTree, SRBTreeNode *pNew) { pNew->left = NULL; pNew->right = NULL; @@ -153,8 +202,12 @@ SRBTreeNode *tRBTreePut(SRBTree *pTree, SRBTreeNode *pNew) { void tRBTreeDrop(SRBTree *pTree, SRBTreeNode *pNode) { // update min/max node - if (pTree->minNode == pNode) pTree->minNode = pNode->parent; - if (pTree->maxNode == pNode) pTree->maxNode = pNode->parent; + if (pTree->minNode == pNode) { + pTree->minNode = tRBTreeSuccessor(pTree->minNode); + } + if (pTree->maxNode == pNode) { + pTree->maxNode = tRBTreePredecessor(pTree->maxNode); + } // drop impl if (pNode->left == NULL) { @@ -208,19 +261,7 @@ void tRBTreeDrop(SRBTree *pTree, SRBTreeNode *pNode) { } SRBTreeNode *tRBTreeDropByKey(SRBTree *pTree, void *pKey) { - SRBTreeNode *pNode = pTree->rootNode; - - while (pNode) { - int32_t c = pTree->cmprFn(pKey, pNode->payload); - - if (c < 0) { - pNode = pNode->left; - } else if (c > 0) { - pNode = pNode->right; - } else { - break; - } - } + SRBTreeNode *pNode = tRBTreeGet(pTree, pKey); if (pNode) { tRBTreeDrop(pTree, pNode); @@ -250,53 +291,14 @@ SRBTreeNode *tRBTreeGet(SRBTree *pTree, void *pKey) { // SRBTreeIter ================================================ SRBTreeNode *tRBTreeIterNext(SRBTreeIter *pIter) { SRBTreeNode *pNode = pIter->pNode; - SRBTree *pTree = pIter->pTree; if (pIter->pNode) { - if (pIter->des) { - // descend - if (pIter->pNode->left) { - pIter->pNode = pIter->pNode->left; - while (pIter->pNode->right) { - pIter->pNode = pIter->pNode->right; - } - } else { - while (true) { - if (pIter->pNode->parent) { - if (pIter->pNode == pIter->pNode->parent->right) { - pIter->pNode = pIter->pNode->parent; - break; - } else { - pIter->pNode = pIter->pNode->parent; - } - } else { - pIter->pNode = NULL; - break; - } - } - } - } else { + if (pIter->asc) { // ascend - if (pIter->pNode->right) { - pIter->pNode = pIter->pNode->right; - while (pIter->pNode->left) { - pIter->pNode = pIter->pNode->left; - } - } else { - while (true) { - if (pIter->pNode->parent) { - if (pIter->pNode == pIter->pNode->parent->left) { - pIter->pNode = pIter->pNode->parent; - break; - } else { - pIter->pNode = pIter->pNode->parent; - } - } else { - pIter->pNode = NULL; - break; - } - } - } + pIter->pNode = tRBTreeSuccessor(pIter->pNode); + } else { + // descend + pIter->pNode = tRBTreePredecessor(pIter->pNode); } } diff --git a/source/util/test/trbtreeTest.cpp b/source/util/test/trbtreeTest.cpp index be62740949a4c4889c45679e4a03fc21927b6a3d..1e0768c32f017cea055eac1bc1ba9dde9a222a4c 100644 --- a/source/util/test/trbtreeTest.cpp +++ b/source/util/test/trbtreeTest.cpp @@ -25,12 +25,13 @@ TEST(trbtreeTest, rbtree_test1) { tRBTreePut(&rt, pNode); } - SRBTreeIter rti = tRBTreeIterCreate(&rt, 0); + SRBTreeIter rti = tRBTreeIterCreate(&rt, 1); SRBTreeNode *pNode = tRBTreeIterNext(&rti); int la = 0; while (pNode) { GTEST_ASSERT_GT(*(int *)pNode->payload, la); la = *(int *)pNode->payload; + // printf("%d\n", la); pNode = tRBTreeIterNext(&rti); } } \ No newline at end of file