10909 - Lucky Number

[ImLazy]

AC now.
The first time I got MLE. After I decreased the array's length I got AC. (There are only 136412 lucky numbers less than 2000000.)

[shanto86]

I am getting TLE using BST. can any one please tell me how to speed it up?

Code: Select all

/*RED BLACK TREE. VERSION:1*/
#include<stdio.h>
#include<algorithm>
using namespace std;
#define max 2000001 //highest size of BST
#define RED 1
#define BLACK 0
/*SAME VALUE CAN BE INSERTED. IN THAT CASE BOTH WILL BE HANDLED AS IDENTICAL*/
struct RB_BST
{
	int val,index,lchild,rchild,parent,color;
	int size;
}NODE[1000001];

int avail=1;
int root,size;
int flag[150000];
int now;

void INIT();
void INSERT(RB_BST);
void INSERT_FIX(RB_BST);
int TREE_SUCCESSOR(RB_BST);
void DELETE(RB_BST);
void DELETE_FIX(RB_BST);
void LEFT_ROTATE(RB_BST);
void RIGHT_ROTATE(RB_BST);
int SEARCH_SERIAL(int);
int SEARCH_VAL(int);

void INIT()
{
	root=0;
	size=0;

	//NULL POINT 0
	NODE[0].val=NODE[0].index=NODE[0].lchild=NODE[0].rchild=NODE[0].parent=NODE[0].color=NODE[0].size=0;
}

void INSERT(RB_BST T)
{
	int a,x,y;

	T.size=1;

	a=avail;
	avail++;
	size++;

	y=0;
	x=root;

	while(x!=0)
	{
		NODE[x].size++;

		y=x;

		if(T.val < NODE[x].val)
			x=NODE[x].lchild;
		else
			x=NODE[x].rchild;
	}

	T.parent=y;
	T.index=a;

	if(y==0)
		root=a;
	else if(T.val < NODE[y].val)
		NODE[y].lchild=a;
	else
		NODE[y].rchild=a;

	T.lchild=0;
	T.rchild=0;
	T.color=RED;

	NODE[a]=T;

	INSERT_FIX(NODE[a]);
}

void INSERT_FIX(RB_BST T)
{
	int y;
	while(NODE[T.parent].color==RED)
	{
		if(T.parent == NODE[NODE[T.parent].parent].lchild )
		{
			y=NODE[NODE[T.parent].parent].rchild;

			if(NODE[y].color==RED)
			{
				NODE[T.parent].color=BLACK;
				NODE[y].color=BLACK;
				NODE[NODE[T.parent].parent].color=RED;
				
				NODE[T.index]=T;
				T=NODE[NODE[T.parent].parent];
			}
			else
			{
				if(T.index == NODE[T.parent].rchild)
				{
					T=NODE[T.parent];
					NODE[T.index]=T;

					LEFT_ROTATE(T);
					T=NODE[T.index];
				}

				NODE[T.parent].color=BLACK;
				NODE[NODE[T.parent].parent].color=RED;

				RIGHT_ROTATE(NODE[NODE[T.parent].parent]);
				T=NODE[T.index];
			}
		}
		else
		{
			y=NODE[NODE[T.parent].parent].lchild;

			if(NODE[y].color==RED)
			{
				NODE[T.parent].color=BLACK;
				NODE[y].color=BLACK;
				NODE[NODE[T.parent].parent].color=RED;
				
				NODE[T.index]=T;
				T=NODE[NODE[T.parent].parent];
			}
			else
			{
				if(T.index == NODE[T.parent].lchild)
				{
					NODE[T.index]=T;
					T=NODE[T.parent];

					RIGHT_ROTATE(T);
					T=NODE[T.index];
				}

				NODE[T.parent].color=BLACK;
				NODE[NODE[T.parent].parent].color=RED;

				LEFT_ROTATE(NODE[NODE[T.parent].parent]);
				T=NODE[T.index];
			}
		}
	}

	NODE[root].color=BLACK;
}

int TREE_SUCCESSOR(RB_BST T)
{
	int x=T.rchild;

	while(NODE[x].lchild!=0)
		x=NODE[x].lchild;

	return x;
}

void DELETE(RB_BST T)
{
	int y,x;

	if(T.lchild==0 || T.rchild==0)
		y=T.index;
	else
		y=TREE_SUCCESSOR(T);

	if(NODE[y].lchild!=0)
		x=NODE[y].lchild;
	else
		x=NODE[y].rchild;

	NODE[x].parent=NODE[y].parent;

	if(NODE[y].parent==0)
		root=x;
	else if(y==NODE[NODE[y].parent].lchild)
		NODE[NODE[y].parent].lchild=x;
	else
		NODE[NODE[y].parent].rchild=x;

	if(y!=T.index)
	{
		NODE[T.index].val=NODE[y].val;
	}

	int z=y;

	while(z!=0)
	{
		NODE[z].size--;
		z=NODE[z].parent;
	}

	if(NODE[y].color==BLACK)
		DELETE_FIX(NODE[x]);

	size--;
	NODE[0].parent=0;
}

void DELETE_FIX(RB_BST T)
{
	int x,w;

	x=T.index;

	while(x!=root && NODE[x].color==BLACK)
	{
		if(x == NODE[NODE[x].parent].lchild)
		{
			w=NODE[NODE[x].parent].rchild;

			if(NODE[w].color==RED)
			{
				NODE[w].color=BLACK;
				NODE[NODE[x].parent].color=RED;
				LEFT_ROTATE(NODE[NODE[x].parent]);
				w=NODE[NODE[x].parent].rchild;
			}

			if(NODE[NODE[w].lchild].color==BLACK && NODE[NODE[w].rchild].color==BLACK)
			{
				NODE[w].color=RED;
				x=NODE[x].parent;
			}
			else
			{
				if(NODE[NODE[w].rchild].color==BLACK)
				{
					NODE[NODE[w].lchild].color=BLACK;
					NODE[w].color=RED;
					RIGHT_ROTATE(NODE[w]);
					w=NODE[NODE[x].parent].rchild;
				}

				NODE[w].color=NODE[NODE[x].parent].color;
				NODE[NODE[x].parent].color=BLACK;
				NODE[NODE[w].rchild].color=BLACK;
				LEFT_ROTATE(NODE[NODE[x].parent]);
				x=root;
			}
		}
		else
		{
			w=NODE[NODE[x].parent].lchild;

			if(NODE[w].color==RED)
			{
				NODE[w].color=BLACK;
				NODE[NODE[x].parent].color=RED;
				RIGHT_ROTATE(NODE[NODE[x].parent]);
				w=NODE[NODE[x].parent].lchild;
			}

			if(NODE[NODE[w].rchild].color==BLACK && NODE[NODE[w].lchild].color==BLACK)
			{
				NODE[w].color=RED;
				x=NODE[x].parent;
			}
			else
			{
				if(NODE[NODE[w].lchild].color==BLACK)
				{
					NODE[NODE[w].rchild].color=BLACK;
					NODE[w].color=RED;
					LEFT_ROTATE(NODE[w]);
					w=NODE[NODE[x].parent].lchild;
				}

				NODE[w].color=NODE[NODE[x].parent].color;
				NODE[NODE[x].parent].color=BLACK;
				NODE[NODE[w].lchild].color=BLACK;
				RIGHT_ROTATE(NODE[NODE[x].parent]);
				x=root;
			}
		}
		NODE[0].parent=0;
	}
	NODE[x].color=BLACK;
}

void LEFT_ROTATE(RB_BST T)
{
	int y;

	y=T.rchild;

	T.rchild=NODE[y].lchild;
	if(NODE[y].lchild!=0) NODE[NODE[y].lchild].parent=T.index;
	NODE[y].parent=NODE[T.parent].index;

	if(T.parent==0)
	{
		root=y;
	}
	else if(T.index == NODE[T.parent].lchild)
	{
		NODE[T.parent].lchild = y;
	}
	else
	{
		NODE[T.parent].rchild = y;
	}

	NODE[y].lchild=T.index;
	T.parent=y;

	NODE[y].size=T.size;
	T.size=NODE[T.lchild].size+NODE[T.rchild].size+1;
	NODE[T.index]=T;
}

void RIGHT_ROTATE(RB_BST T)
{
	int y;

	y=T.lchild;

	T.lchild=NODE[y].rchild;
	if(NODE[y].rchild!=0) NODE[NODE[y].rchild].parent=T.index;
	NODE[y].parent=NODE[T.parent].index;

	if(T.parent==0)
	{
		root=y;
	}
	else if(T.index == NODE[T.parent].rchild)
	{
		NODE[T.parent].rchild = y;
	}
	else
	{
		NODE[T.parent].lchild = y;
	}

	NODE[y].rchild=T.index;
	T.parent=y;

	NODE[y].size=T.size;
	T.size=NODE[T.lchild].size+NODE[T.rchild].size+1;
	NODE[T.index]=T;
}

int SEARCH_VAL(int k)
{
	int x=root;

	while(NODE[x].val != k && x>0)
	{
		if(k < NODE[x].val)
			x=NODE[x].lchild;
		else
			x=NODE[x].rchild;
	}

	if(x==0)
		return -1;
	return x;
}

int SEARCH_SERIAL(int k)
{
	int x;

	if(k > size)
		return -1;

	x=root;

	while(NODE[NODE[x].lchild].size+1!=k)
	{
		if(k <= NODE[NODE[x].lchild].size)
			x=NODE[x].lchild;
		else
		{
			k=k-NODE[NODE[x].lchild].size-1;
			x=NODE[x].rchild;
		}
	}
	return x;
}

/////////////////////////////////////////////////////////////////////////////////////////////

void WORK()
{
	int n,OK;
	int *a,*b;

	while(scanf("%d",&n)!=EOF)
	{
		a=lower_bound(flag,flag+now+1,int(n/2));
		b=a;

		OK=0;

		while(a-flag>=0 && b-flag<=now)
		{
			if((*a)+(*b)==n)
			{
				OK=1;
				break;
			}
			else if((*a)+(*b)>n)
				a--;
			else
				b++;
		}

		if(OK)
			printf("%d is the sum of %d and %d.\n",n,*a,*b);
		else
			printf("%d is not the sum of two luckies!\n",n);
	}
}

	
int main()
{
	int i,j,k;

	RB_BST temp;

	INIT();

	for(i=1;i<max;i+=2)
	{
		temp.val=i;
		INSERT(temp);
	}

	now=0;
	flag[0]=1;
	for(k=2;k<=size;k++)
	{
		i=SEARCH_SERIAL(k);
		i=NODE[i].val;
		flag[++now]=i;

		for(j=i;j<=size;j+=i-1)
		{
			DELETE(NODE[SEARCH_SERIAL(j)]);
		}
	}

	printf("%d\n",size);

	WORK();

	return 0;
}

[inproblem]

Hi,
I m a newbie. will u pls tell me how to solve the problem
with BST ? How i'll construct the tree initially? How the child-parent
elements will be chosen? What informations will be kept in each node?

Mainly i want to know , which elements will be parents nd which will
be children? How shall i remove k'th element each time?

[fpavetic]

firstly, think about the order you should insert elements to keep the tree balanced;

node should keep two informations: key and size ( number of elements in its left subtree + right subtree + 1 ); if you use an array and node index is i then its left subtree index is 2*i+1 and right 2*i+2;
as for deletion, just mark the key on some dummy value and decrease size member for that node and all of its parents.

to find k-th element make use of size member in a node. think about when k-th element will be in left subtree, right subtree, or when it is a root