#include <math.h>
#include <stdlib.h>

typedef struct RedBlackTreeNode
{
	int							key;
	long double					count;
	int							color;
	struct	RedBlackTreeNode*	left;
	struct	RedBlackTreeNode*	right;
}rbtree;

static rbtree* RedBlackTree_insert(rbtree *p, int key, long double count);
static rbtree* RedBlackTree_find(rbtree *p, int key);
static void RedBlackTree_free(rbtree *p);
long double partition_numbers_imp(int n, rbtree **rootp);

long double partition_numbers(int n)
{
	long double count;
	rbtree *root = NULL;

	count = partition_numbers_imp(n, &root);
	RedBlackTree_free(root);
	return count;
}

long double partition_numbers_imp(int n, rbtree **rootp)
{
	int k;
	long double count;
	rbtree *p;

	if(n < 0)
		return 0.0;
	if(n == 0)
		return 1.0;

	/* if the count for n is already added in tree, return the count directly from tree */
	if((p = RedBlackTree_find(*rootp, n)) != NULL)
		return p->count;

	for(k = 1, count = 0; k <= n; k++)
		count += pow((long double)(-1), k+1) * (partition_numbers_imp(n - ((k*(3*k-1))/2), rootp) + partition_numbers_imp(n - ((k*(3*k+1))/2), rootp));

	/* remember the the count for n in redblack tree */
	if((p = RedBlackTree_find(*rootp, n)) == NULL)
		*rootp = RedBlackTree_insert(*rootp, n, count);
	return count;
}



static rbtree* RedBlackTree_insert_imp(rbtree *p, int key, long double count, int direction);
static rbtree* rightRotation(rbtree *p);
static rbtree* leftRotation(rbtree *p);
static rbtree* create_node(int key, long double count);

enum {BLACK = 0, RED};
enum {LEFT = 0, RIGHT};

/* insert function */
static rbtree* RedBlackTree_insert(rbtree *p, int key, long double count)
{
	rbtree *root;

	root = RedBlackTree_insert_imp(p, key, count, LEFT);

	/* the color of root should always be black */
	root->color = BLACK;
	return root;
}

/* lookup function */
static rbtree* RedBlackTree_find(rbtree *p, int key)
{
	while(p)
	{
		if(p->key == key)
			return p;
		else if(key < p->key)
			p = p->left;
		else
			p = p->right;
	}
	return NULL;
}

static void RedBlackTree_free(rbtree *p)
{
	if(p == NULL)
		return;

	RedBlackTree_free(p->left);
	RedBlackTree_free(p->right);
	free(p);
}

static rbtree* RedBlackTree_insert_imp(rbtree *p, int key, long double count, int direction)
{
	if(p == NULL)
		return create_node(key, count);

	/* while moving from top to bottom, check for 4-nodes and split them*/
	if(p->left != NULL && p->right != NULL)
	{
		if((p->left->color == RED) && (p->right->color == RED))
		{
			/* split the 4-node into two 2-nodes by changing the color of left and right
				child from red to black. Also, convert the parent of 4-node from n-node to (n+1)-node */
			p->color = RED;
			p->left->color = BLACK;
			p->right->color = BLACK;
		}
	}

	if(key < p->key)
	{
		/* move to left subtree */
		p->left = RedBlackTree_insert_imp(p->left, key, count, LEFT);

		/* three nodes are connected by red links. the nodes orientation is n1--right--p--left-n3.
			rotate p right */
		if(p->color == RED			&&
			p->left->color == RED	&&
			direction == RIGHT)
		{
			p = rightRotation(p);
		}

		/* three nodes are connected by red links. the nodes orientation is n1--left--p--left-n3.
			rotate p right */
		if(p->left != NULL			&&
			p->left->left != NULL	&&
			p->left->color == RED	&&
			p->left->left->color == RED)
		{
			p = rightRotation(p);

			/* change the color of p to black and right child to red */
			p->color = BLACK;
			p->right->color = RED;
		}
	}
	else
	{
		/* move to right subtree */
		p->right = RedBlackTree_insert_imp(p->right, key, count, RIGHT);

		/* three nodes are connected by red links. the nodes orientation is n1--left--p--right-n3.
			rotate p left */
		if(p->color == RED			&&
			p->right->color == RED	&&
			direction == LEFT)
		{
			p = leftRotation(p);
		}

		/* three nodes are connected by red links. the nodes orientation is n1--right--p--right-n3.
			rotate p left */
		if(p->right != NULL			&&
			p->right->right != NULL	&&
			p->right->color == RED	&&
			p->right->right->color == RED)
		{
			p = leftRotation(p);

			/* change the color of p to black and right child to red */
			p->color = BLACK;
			p->left->color = RED;
		}
	}

	return p;
}

static rbtree* create_node(int key, long double count)
{
	rbtree *p;

	p = (rbtree*)malloc(sizeof(*p));

	if(p == NULL)
		return NULL;

	p->color = RED;
	p->key = key;
	p->count = count;
	p->left = p->right = NULL;
	return p;
}

static rbtree* rightRotation(rbtree *p)
{
	rbtree *ptr = p->left;

	p->left = ptr->right;
	ptr->right = p;

	return ptr;
}

static rbtree* leftRotation(rbtree *p)
{
	rbtree *ptr = p->right;

	p->right = ptr->left;
	ptr->left = p;

	return ptr;
}