TopCoder SRM 473 Div 1 Medium RightTriangle

本番でテスト落ちたやつ。

問題概要

円周上をn等分する点に、時計回りに0,1,2,...,n-1と番号をつける。
そして、0,1,2,...,m-1のそれぞれのkに対して、
l=(a*a*k+b*k+c)%nとし、lから数えて初めて最初に赤でない点を赤にする。

この操作を終えたとき、円周上の異なる3つの赤い点を結んでできる直角三角形はいくつできるかを求めよ。
n≦1000000, 赤い点の数mはm≦100000を満たす。

解法

赤い点を作るのにまだ赤くなっていない点を探す必要があるが、これを線形探索で行うと明らかにTLE.
なので平方分割した。√nの大きさの区間にnを区切り、それぞれに赤い点を何個作ったかを数えておく。

本番では次の区間に移動するときに区間の最初からではなく途中から調べていてテストで落ちたorz

ソースコード

int P,sqp,last;
int cnt[1000];
bool red[1000000];

void ins(ll n)
{
	int m,lim;
	while(1)
	{
		m=n/sqp;
		while(cnt[m]==sqp||m==sqp-1&&cnt[m]==last)
		{
			m++; n=m*sqp;
			if(n>=P)n=m=0;
		}
		
		lim=min(sqp*(m+1),P);
		
		while(n<lim&&red[n])n++;
		if(n<lim)
		{
			red[n]=1;
			cnt[m]++; return;
		}
		
		if(n==P)n=0;
	}
}

class RightTriangle {
	public:
	long long triangleCount(int places, int points, int A, int B, int C) 
	{
		P=places; sqp=(int)ceil(sqrt(P*1.0));
		last=P%sqp; if(last==0)last=sqp;
		
		dbg(P);dbg(sqp);dbg(last);
		
		if(P%2)return 0;
		fill_n(cnt,1000,0);
		fill_n(red,P,0);
		
		ll a=A,b=B,c=C;
		rep(i,points)
		{
			ll n=(a*i*i+b*i+c)%P;
			ins(n);
		}
		
		ll ans=0;
		rep(i,P/2)if(red[i])
		{
			if(red[i+P/2])
			ans+=points-2;
		}
		return ans;
	}
};