/*
* ctrees.c
*
* Author: mozman
* Copyright (c) 2010-2013 by Manfred Moitzi
* License: MIT-License
*/
#include "ctrees.h"
#include "stack.h"
#include <Python.h>
#define LEFT 0
#define RIGHT 1
#define KEY(node) (node->key)
#define VALUE(node) (node->value)
#define LEFT_NODE(node) (node->link[LEFT])
#define RIGHT_NODE(node) (node->link[RIGHT])
#define LINK(node, dir) (node->link[dir])
#define XDATA(node) (node->xdata)
#define RED(node) (node->xdata)
#define BALANCE(node) (node->xdata)
static node_t *
ct_new_node(PyObject *key, PyObject *value, int xdata)
{
node_t *new_node = PyMem_Malloc(sizeof(node_t));
if (new_node != NULL) {
KEY(new_node) = key;
Py_INCREF(key);
VALUE(new_node) = value;
Py_INCREF(value);
LEFT_NODE(new_node) = NULL;
RIGHT_NODE(new_node) = NULL;
XDATA(new_node) = xdata;
}
return new_node;
}
static void
ct_delete_node(node_t *node)
{
if (node != NULL) {
Py_XDECREF(KEY(node));
Py_XDECREF(VALUE(node));
LEFT_NODE(node) = NULL;
RIGHT_NODE(node) = NULL;
PyMem_Free(node);
}
}
extern void
ct_delete_tree(node_t *root)
{
if (root == NULL)
return;
if (LEFT_NODE(root) != NULL) {
ct_delete_tree(LEFT_NODE(root));
}
if (RIGHT_NODE(root) != NULL) {
ct_delete_tree(RIGHT_NODE(root));
}
ct_delete_node(root);
}
static void
ct_swap_data(node_t *node1, node_t *node2)
{
PyObject *tmp;
tmp = KEY(node1);
KEY(node1) = KEY(node2);
KEY(node2) = tmp;
tmp = VALUE(node1);
VALUE(node1) = VALUE(node2);
VALUE(node2) = tmp;
}
int
ct_compare(PyObject *key1, PyObject *key2)
{
int res;
res = PyObject_RichCompareBool(key1, key2, Py_LT);
if (res > 0)
return -1;
else if (res < 0) {
PyErr_SetString(PyExc_TypeError, "invalid type for key");
return 0;
}
/* second compare:
+1 if key1 > key2
0 if not -> equal
-1 means error, if error, it should happend at the first compare
*/
return PyObject_RichCompareBool(key1, key2, Py_GT);
}
extern node_t *
ct_find_node(node_t *root, PyObject *key)
{
int res;
while (root != NULL) {
res = ct_compare(key, KEY(root));
if (res == 0) /* key found */
return root;
else {
root = LINK(root, (res > 0));
}
}
return NULL; /* key not found */
}
extern PyObject *
ct_get_item(node_t *root, PyObject *key)
{
node_t *node;
PyObject *tuple;
node = ct_find_node(root, key);
if (node != NULL) {
tuple = PyTuple_New(2);
PyTuple_SET_ITEM(tuple, 0, KEY(node));
PyTuple_SET_ITEM(tuple, 1, VALUE(node));
return tuple;
}
Py_RETURN_NONE;
}
extern node_t *
ct_max_node(node_t *root)
/* get node with largest key */
{
if (root == NULL)
return NULL;
while (RIGHT_NODE(root) != NULL)
root = RIGHT_NODE(root);
return root;
}
extern node_t *
ct_min_node(node_t *root)
// get node with smallest key
{
if (root == NULL)
return NULL;
while (LEFT_NODE(root) != NULL)
root = LEFT_NODE(root);
return root;
}
extern int
ct_bintree_remove(node_t **rootaddr, PyObject *key)
/* attention: rootaddr is the address of the root pointer */
{
node_t *node, *parent, *replacement;
int direction, cmp_res, down_dir;
node = *rootaddr;
if (node == NULL)
return 0; /* root is NULL */
parent = NULL;
direction = 0;
while (1) {
cmp_res = ct_compare(key, KEY(node));
if (cmp_res == 0) /* key found, remove node */
{
if ((LEFT_NODE(node) != NULL) && (RIGHT_NODE(node) != NULL)) {
/* find replacement node: smallest key in right-subtree */
parent = node;
direction = RIGHT;
replacement = RIGHT_NODE(node);
while (LEFT_NODE(replacement) != NULL) {
parent = replacement;
direction = LEFT;
replacement = LEFT_NODE(replacement);
}
LINK(parent, direction) = RIGHT_NODE(replacement);
/* swap places */
ct_swap_data(node, replacement);
node = replacement; /* delete replacement node */
}
else {
down_dir = (LEFT_NODE(node) == NULL) ? RIGHT : LEFT;
if (parent == NULL) /* root */
{
*rootaddr = LINK(node, down_dir);
}
else {
LINK(parent, direction) = LINK(node, down_dir);
}
}
ct_delete_node(node);
return 1; /* remove was success full */
}
else {
direction = (cmp_res < 0) ? LEFT : RIGHT;
parent = node;
node = LINK(node, direction);
if (node == NULL)
return 0; /* error key not found */
}
}
}
extern int
ct_bintree_insert(node_t **rootaddr, PyObject *key, PyObject *value)
/* attention: rootaddr is the address of the root pointer */
{
node_t *parent, *node;
int direction, cval;
node = *rootaddr;
if (node == NULL) {
node = ct_new_node(key, value, 0); /* new node is also the root */
if (node == NULL)
return -1; /* got no memory */
*rootaddr = node;
}
else {
direction = LEFT;
parent = NULL;
while (1) {
if (node == NULL) {
node = ct_new_node(key, value, 0);
if (node == NULL)
return -1; /* get no memory */
LINK(parent, direction) = node;
return 1;
}
cval = ct_compare(key, KEY(node));
if (cval == 0) {
/* key exists, replace value object */
Py_XDECREF(VALUE(node)); /* release old value object */
VALUE(node) = value; /* set new value object */
Py_INCREF(value); /* take new value object */
return 0;
}
else {
parent = node;
direction = (cval < 0) ? LEFT : RIGHT;
node = LINK(node, direction);
}
}
}
return 1;
}
static int
is_red (node_t *node)
{
return (node != NULL) && (RED(node) == 1);
}
#define rb_new_node(key, value) ct_new_node(key, value, 1)
static node_t *
rb_single(node_t *root, int dir)
{
node_t *save = root->link[!dir];
root->link[!dir] = save->link[dir];
save->link[dir] = root;
RED(root) = 1;
RED(save) = 0;
return save;
}
static node_t *
rb_double(node_t *root, int dir)
{
root->link[!dir] = rb_single(root->link[!dir], !dir);
return rb_single(root, dir);
}
#define rb_new_node(key, value) ct_new_node(key, value, 1)
extern int
rb_insert(node_t **rootaddr, PyObject *key, PyObject *value)
{
node_t *root = *rootaddr;
if (root == NULL) {
/*
We have an empty tree; attach the
new node directly to the root
*/
root = rb_new_node(key, value);
if (root == NULL)
return -1; // got no memory
}
else {
node_t head; /* False tree root */
node_t *g, *t; /* Grandparent & parent */
node_t *p, *q; /* Iterator & parent */
int dir = 0;
int last = 0;
int new_node = 0;
/* Set up our helpers */
t = &head;
g = NULL;
p = NULL;
RIGHT_NODE(t) = root;
LEFT_NODE(t) = NULL;
q = RIGHT_NODE(t);
/* Search down the tree for a place to insert */
for (;;) {
int cmp_res;
if (q == NULL) {
/* Insert a new node at the first null link */
q = rb_new_node(key, value);
p->link[dir] = q;
new_node = 1;
if (q == NULL)
return -1; // get no memory
}
else if (is_red(q->link[0]) && is_red(q->link[1])) {
/* Simple red violation: color flip */
RED(q) = 1;
RED(q->link[0]) = 0;
RED(q->link[1]) = 0;
}
if (is_red(q) && is_red(p)) {
/* Hard red violation: rotations necessary */
int dir2 = (t->link[1] == g);
if (q == p->link[last])
t->link[dir2] = rb_single(g, !last);
else
t->link[dir2] = rb_double(g, !last);
}
/* Stop working if we inserted a new node. */
if (new_node)
break;
cmp_res = ct_compare(KEY(q), key);
if (cmp_res == 0) { /* key exists? */
Py_XDECREF(VALUE(q)); /* release old value object */
VALUE(q) = value; /* set new value object */
Py_INCREF(value); /* take new value object */
return 0;
}
last = dir;
dir = (cmp_res < 0);
/* Move the helpers down */
if (g != NULL)
t = g;
g = p;
p = q;
q = q->link[dir];
}
/* Update the root (it may be different) */
root = head.link[1];
}
/* Make the root black for simplified logic */
RED(root) = 0;
(*rootaddr) = root;
return 1;
}
extern int
rb_remove(node_t **rootaddr, PyObject *key)
{
node_t *root = *rootaddr;
node_t head = { { NULL } }; /* False tree root */
node_t *q, *p, *g; /* Helpers */
node_t *f = NULL; /* Found item */
int dir = 1;
if (root == NULL)
return 0;
/* Set up our helpers */
q = &head;
g = p = NULL;
RIGHT_NODE(q) = root;
/*
Search and push a red node down
to fix red violations as we go
*/
while (q->link[dir] != NULL) {
int last = dir;
int cmp_res;
/* Move the helpers down */
g = p, p = q;
q = q->link[dir];
cmp_res = ct_compare(KEY(q), key);
dir = cmp_res < 0;
/*
Save the node with matching data and keep
going; we'll do removal tasks at the end
*/
if (cmp_res == 0)
f = q;
/* Push the red node down with rotations and color flips */
if (!is_red(q) && !is_red(q->link[dir])) {
if (is_red(q->link[!dir]))
p = p->link[last] = rb_single(q, dir);
else if (!is_red(q->link[!dir])) {
node_t *s = p->link[!last];
if (s != NULL) {
if (!is_red(s->link[!last]) &&
!is_red(s->link[last])) {
/* Color flip */
RED(p) = 0;
RED(s) = 1;
RED(q) = 1;
}
else {
int dir2 = g->link[1] == p;
if (is_red(s->link[last]))
g->link[dir2] = rb_double(p, last);
else if (is_red(s->link[!last]))
g->link[dir2] = rb_single(p, last);
/* Ensure correct coloring */
RED(q) = RED(g->link[dir2]) = 1;
RED(g->link[dir2]->link[0]) = 0;
RED(g->link[dir2]->link[1]) = 0;
}
}
}
}
}
/* Replace and remove the saved node */
if (f != NULL) {
ct_swap_data(f, q);
p->link[p->link[1] == q] = q->link[q->link[0] == NULL];
ct_delete_node(q);
}
/* Update the root (it may be different) */
root = head.link[1];
/* Make the root black for simplified logic */
if (root != NULL)
RED(root) = 0;
*rootaddr = root;
return (f != NULL);
}
#define avl_new_node(key, value) ct_new_node(key, value, 0)
#define height(p) ((p) == NULL ? -1 : (p)->xdata)
#define avl_max(a, b) ((a) > (b) ? (a) : (b))
static node_t *
avl_single(node_t *root, int dir)
{
node_t *save = root->link[!dir];
int rlh, rrh, slh;
/* Rotate */
root->link[!dir] = save->link[dir];
save->link[dir] = root;
/* Update balance factors */
rlh = height(root->link[0]);
rrh = height(root->link[1]);
slh = height(save->link[!dir]);
BALANCE(root) = avl_max(rlh, rrh) + 1;
BALANCE(save) = avl_max(slh, BALANCE(root)) + 1;
return save;
}
static node_t *
avl_double(node_t *root, int dir)
{
root->link[!dir] = avl_single(root->link[!dir], !dir);
return avl_single(root, dir);
}
extern int
avl_insert(node_t **rootaddr, PyObject *key, PyObject *value)
{
node_t *root = *rootaddr;
if (root == NULL) {
root = avl_new_node(key, value);
if (root == NULL)
return -1; // got no memory
}
else {
node_t *it, *up[32];
int upd[32], top = 0;
int done = 0;
int cmp_res;
it = root;
/* Search for an empty link, save the path */
for (;;) {
/* Push direction and node onto stack */
cmp_res = ct_compare(KEY(it), key);
if (cmp_res == 0) {
Py_XDECREF(VALUE(it)); // release old value object
VALUE(it) = value; // set new value object
Py_INCREF(value); // take new value object
return 0;
}
// upd[top] = it->data < data;
upd[top] = (cmp_res < 0);
up[top++] = it;
if (it->link[upd[top - 1]] == NULL)
break;
it = it->link[upd[top - 1]];
}
/* Insert a new node at the bottom of the tree */
it->link[upd[top - 1]] = avl_new_node(key, value);
if (it->link[upd[top - 1]] == NULL)
return -1; // got no memory
/* Walk back up the search path */
while (--top >= 0 && !done) {
// int dir = (cmp_res < 0);
int lh, rh, max;
cmp_res = ct_compare(KEY(up[top]), key);
lh = height(up[top]->link[upd[top]]);
rh = height(up[top]->link[!upd[top]]);
/* Terminate or rebalance as necessary */
if (lh - rh == 0)
done = 1;
if (lh - rh >= 2) {
node_t *a = up[top]->link[upd[top]]->link[upd[top]];
node_t *b = up[top]->link[upd[top]]->link[!upd[top]];
if (height( a ) >= height( b ))
up[top] = avl_single(up[top], !upd[top]);
else
up[top] = avl_double(up[top], !upd[top]);
/* Fix parent */
if (top != 0)
up[top - 1]->link[upd[top - 1]] = up[top];
else
root = up[0];
done = 1;
}
/* Update balance factors */
lh = height(up[top]->link[upd[top]]);
rh = height(up[top]->link[!upd[top]]);
max = avl_max(lh, rh);
BALANCE(up[top]) = max + 1;
}
}
(*rootaddr) = root;
return 1;
}
extern int
avl_remove(node_t **rootaddr, PyObject *key)
{
node_t *root = *rootaddr;
int cmp_res;
if (root != NULL) {
node_t *it, *up[32];
int upd[32], top = 0;
it = root;
for (;;) {
/* Terminate if not found */
if (it == NULL)
return 0;
cmp_res = ct_compare(KEY(it), key);
if (cmp_res == 0)
break;
/* Push direction and node onto stack */
upd[top] = (cmp_res < 0);
up[top++] = it;
it = it->link[upd[top - 1]];
}
/* Remove the node */
if (it->link[0] == NULL ||
it->link[1] == NULL) {
/* Which child is not null? */
int dir = it->link[0] == NULL;
/* Fix parent */
if (top != 0)
up[top - 1]->link[upd[top - 1]] = it->link[dir];
else
root = it->link[dir];
ct_delete_node(it);
}
else {
/* Find the inorder successor */
node_t *heir = it->link[1];
/* Save the path */
upd[top] = 1;
up[top++] = it;
while ( heir->link[0] != NULL ) {
upd[top] = 0;
up[top++] = heir;
heir = heir->link[0];
}
/* Swap data */
ct_swap_data(it, heir);
/* Unlink successor and fix parent */
up[top - 1]->link[up[top - 1] == it] = heir->link[1];
ct_delete_node(heir);
}
/* Walk back up the search path */
while (--top >= 0) {
int lh = height(up[top]->link[upd[top]]);
int rh = height(up[top]->link[!upd[top]]);
int max = avl_max(lh, rh);
/* Update balance factors */
BALANCE(up[top]) = max + 1;
/* Terminate or rebalance as necessary */
if (lh - rh == -1)
break;
if (lh - rh <= -2) {
node_t *a = up[top]->link[!upd[top]]->link[upd[top]];
node_t *b = up[top]->link[!upd[top]]->link[!upd[top]];
if (height(a) <= height(b))
up[top] = avl_single(up[top], upd[top]);
else
up[top] = avl_double(up[top], upd[top]);
/* Fix parent */
if (top != 0)
up[top - 1]->link[upd[top - 1]] = up[top];
else
root = up[0];
}
}
}
(*rootaddr) = root;
return 1;
}
extern node_t *
ct_succ_node(node_t *root, PyObject *key)
{
node_t *succ = NULL;
node_t *node = root;
int cval;
while (node != NULL) {
cval = ct_compare(key, KEY(node));
if (cval == 0)
break;
else if (cval < 0) {
if ((succ == NULL) ||
(ct_compare(KEY(node), KEY(succ)) < 0))
succ = node;
node = LEFT_NODE(node);
} else
node = RIGHT_NODE(node);
}
if (node == NULL)
return NULL;
/* found node of key */
if (RIGHT_NODE(node) != NULL) {
/* find smallest node of right subtree */
node = RIGHT_NODE(node);
while (LEFT_NODE(node) != NULL)
node = LEFT_NODE(node);
if (succ == NULL)
succ = node;
else if (ct_compare(KEY(node), KEY(succ)) < 0)
succ = node;
}
return succ;
}
extern node_t *
ct_prev_node(node_t *root, PyObject *key)
{
node_t *prev = NULL;
node_t *node = root;
int cval;
while (node != NULL) {
cval = ct_compare(key, KEY(node));
if (cval == 0)
break;
else if (cval < 0)
node = LEFT_NODE(node);
else {
if ((prev == NULL) || (ct_compare(KEY(node), KEY(prev)) > 0))
prev = node;
node = RIGHT_NODE(node);
}
}
if (node == NULL) /* stay at dead end (None) */
return NULL;
/* found node of key */
if (LEFT_NODE(node) != NULL) {
/* find biggest node of left subtree */
node = LEFT_NODE(node);
while (RIGHT_NODE(node) != NULL)
node = RIGHT_NODE(node);
if (prev == NULL)
prev = node;
else if (ct_compare(KEY(node), KEY(prev)) > 0)
prev = node;
}
return prev;
}
extern node_t *
ct_floor_node(node_t *root, PyObject *key)
{
node_t *prev = NULL;
node_t *node = root;
int cval;
while (node != NULL) {
cval = ct_compare(key, KEY(node));
if (cval == 0)
return node;
else if (cval < 0)
node = LEFT_NODE(node);
else {
if ((prev == NULL) || (ct_compare(KEY(node), KEY(prev)) > 0))
prev = node;
node = RIGHT_NODE(node);
}
}
return prev;
}
extern node_t *
ct_ceiling_node(node_t *root, PyObject *key)
{
node_t *succ = NULL;
node_t *node = root;
int cval;
while (node != NULL) {
cval = ct_compare(key, KEY(node));
if (cval == 0)
return node;
else if (cval < 0) {
if ((succ == NULL) ||
(ct_compare(KEY(node), KEY(succ)) < 0))
succ = node;
node = LEFT_NODE(node);
} else
node = RIGHT_NODE(node);
}
return succ;
}
extern int
ct_index_of(node_t *root, PyObject *key)
/*
get index of item <key>, returns -1 if key not found.
*/
{
node_t *node = root;
int index = 0;
int go_down = 1;
node_stack_t *stack;
stack = stack_init(32);
for (;;) {
if ((LEFT_NODE(node) != NULL) && go_down) {
stack_push(stack, node);
node = LEFT_NODE(node);
}
else {
if (ct_compare(KEY(node), key) == 0) {
stack_delete(stack);
return index;
}
index++;
if (RIGHT_NODE(node) != NULL) {
node = RIGHT_NODE(node);
go_down = 1;
}
else {
if (stack_is_empty(stack)) {
stack_delete(stack);
return -1;
}
node = stack_pop(stack);
go_down = 0;
}
}
}
}
extern node_t *
ct_node_at(node_t *root, int index)
{
/*
root -- root node of tree
index -- index of wanted node
return NULL if index out of range
*/
node_t *node = root;
int counter = 0;
int go_down = 1;
node_stack_t *stack;
if (index < 0) return NULL;
stack = stack_init(32);
for(;;) {
if ((LEFT_NODE(node) != NULL) && go_down) {
stack_push(stack, node);
node = LEFT_NODE(node);
}
else {
if (counter == index) {
/* reached wanted index */
stack_delete(stack);
return node;
}
counter++;
if (RIGHT_NODE(node) != NULL) {
node = RIGHT_NODE(node);
go_down = 1;
}
else {
if (stack_is_empty(stack)) { /* index out of range */
stack_delete(stack);
return NULL;
}
node = stack_pop(stack);
go_down = 0;
}
}
}
}