mediantree.c

#include <assert.h>
#include <limits.h>
#include <stdio.h>
#include <stdlib.h>

struct node {
	struct node *left, *right, *parent;
	int value;
};

struct tree {
	struct node *root;
	struct node *median;
	int offset;
};

#define ISTHREAD(pointer) ((long)pointer & 1)
#define SETTHREAD(pointer) ((struct node *)((long)pointer | 1))
#define GETPOINTER(thread) ((struct node *)((long)thread & ~0x1))

void
mediantree_addnode(struct node *root, struct node *node)
{

	if (node->value > root->value) {
		if (!ISTHREAD(root->right)) {
			mediantree_addnode(root->right, node);
		} else {
			struct node *tmp = root->right;
			root->right = node;
			node->right = tmp;
			node->left = SETTHREAD(root);
			node->parent = root;
		}
	}
	else {
		if (!ISTHREAD(root->left)) {
			mediantree_addnode(root->left, node);
		} else {
			struct node *tmp = root->left;
			root->left = node;
			node->left = tmp;
			node->right = SETTHREAD(root);
			node->parent = root;
		}
	}
}

struct node *
mediantree_smallest(struct node *node)
{
	while (!ISTHREAD(node->left))
		node = node->left;

	return node;
}

struct node *
mediantree_biggest(struct node *node)
{
	while (!ISTHREAD(node->right))
		node = node->right;

	return node;
}

struct node *
mediantree_add(struct tree *tree, int value)
{
	struct node *node = calloc(1, sizeof(struct node));
	node->value = value;

	mediantree_addnode(tree->root, node);
	if (value <= tree->median->value)
		tree->offset--;
	else
		tree->offset++;

	if (tree->offset < 0) {
		if (ISTHREAD(tree->median->left))
			tree->median = GETPOINTER(tree->median->left);
		else
			tree->median = mediantree_biggest(tree->median->left);
		tree->offset += 2;
	} else if (tree->offset > 1) {
		if (ISTHREAD(tree->median->right))
			tree->median = GETPOINTER(tree->median->right);
		else
			tree->median = mediantree_smallest(tree->median->right);

		tree->offset -= 2;
	}

	return node;
}

float
mediantree_median(struct tree *tree)
{
	if (tree->offset == 1)
		return (tree->median->value +
			GETPOINTER(tree->median->right)->value) / 2.0;

	return tree->median->value;
}

struct tree *
mediantree_new(void)
{
	struct tree *tree = calloc(1, sizeof(struct tree));
	tree->root = calloc(1, sizeof(struct node));

	tree->median = tree->root;
	tree->root->value = INT_MAX;
	tree->root->right = tree->root;
	tree->root->left = SETTHREAD(tree->root);
	tree->root->parent = NULL;
	tree->offset = -1;

	return tree;
}

int
main(void)
{
	struct tree *tree = mediantree_new();
	struct node *node;

	assert(mediantree_median(tree) == INT_MAX);
	mediantree_add(tree, 5);
	mediantree_add(tree, 10);
	mediantree_add(tree, 1);
	mediantree_add(tree, 15);
	assert(mediantree_median(tree) == 7.5);
	node = mediantree_add(tree, 3);
	assert(mediantree_median(tree) == 5);
	mediantree_add(tree, 4);
	assert(mediantree_median(tree) == 4.5);
	mediantree_add(tree, 7);
	mediantree_add(tree, 11);
	mediantree_add(tree, 16);
	assert(mediantree_median(tree) == 7);
	mediantree_add(tree, 17);
	mediantree_add(tree, 18);
	mediantree_add(tree, 19);
	mediantree_add(tree, 20);
	assert(mediantree_median(tree) == 11);
	mediantree_add(tree, 0);
	mediantree_add(tree, -1);
	mediantree_add(tree, -5);
	mediantree_add(tree, -4);
	assert(mediantree_median(tree) == 7);
	mediantree_add(tree, 7);
	mediantree_add(tree, 7);
	mediantree_add(tree, 7);
	assert(mediantree_median(tree) == 7);
	mediantree_add(tree, 25);
	mediantree_add(tree, 50);
	mediantree_add(tree, 75);
	mediantree_add(tree, 100);
	assert(mediantree_median(tree) == 8.5);

	return 0;
}