TopCoder SRM 327 Div1 Medium QuadraticEquations

問題

tについての方程式at^2 + bt + c = 0であって、
係数a, b, cが-k以上k以下の整数で、
t = (x + y * √d) / zを解の一つとして持つようなものはいくつあるか、求めよ。

制約条件

0≦k≦10^6
-1000≦x, y, z≦1000
zは0でない
1≦d≦1000

方針

dが平方数でない場合とそうでない場合で場合わけする。

dが平方数でないとき、tの解はt = (x ± y * √d) / zになる。

このとき、a, b, cはx, y, d, zの式で表せる。
a, b, cの公約数を1にしておいて、a, b, cのうち絶対値が最も小さいもの分だけ
方程式に定数倍の自由度があるので、それが答え。

dが平方数のとき、tの解はt = p / q（p, qはx, y, z, dから計算される値）と書ける。
もう一つの解をt = u / vと置く。
すると、方程式は
qut^2 - (pu + qv)t + pv = 0となる。

aについての制約から、
-k≦qu≦k ⇔ -k / q ≦ u ≦ k / q
なので、この範囲でuを一つずつループを回して決めうちする。

bについての制約から、
-k≦pv≦k
cについての制約から、
(-k + pu)/q ≦ v ≦ (k + pu)/q

がそれぞれ成り立つので、vに関する連立不等式となり、
可能なvの値の範囲がわかるので、このときの式の個数がわかる。
これを全ての可能なuについて足し合わせればいい。

ソースコード

// solve(a / c <= x <= b / c)
pi solve(ll a, ll b, ll c){
	if(c < 0) return solve(-b, -a, -c);
	if(c == 0){
		if(a <= 0 && b >= 0) return mp(-inf, inf);
		return mp(inf, -inf);
	}
	
	ll s = a / c - 2, t = b / c + 2;
	while(s * c < a) s++;
	while(t * c > b) t--;
	return mp(s, t);
}

class QuadraticEquations {
	public:
	long long howMany(int x, int y, int d, int z, int k) 
	{
		if(y == 0) d = 1;
		
		int sd = (int)sqrt(d + 1e-9);
		if(sd * sd == d){
			ll q = z, p = x + y * sd;
			ll g = __gcd(q, p);
			q /= g; p /= g;
			if(q < 0) q *= -1, p *= -1;
			
			ll ans = 0;
			for(int u = -k; u <= k; u++){
				
				if(q * u < -k || q * u > k) continue;
				
				pi a = solve(-k + p * u, k + p * u, q);
				pi b = solve(-k, k, p);
				
				ll lo = max(a.first, b.first), hi = min(a.second, b.second);
				if(lo <= hi) ans += hi - lo + 1;
			}
			return ans;
		}
		else{
			ll g = __gcd(__gcd((ll)z * z, 2ll * x * z), x * x - (ll)y * y * d);
			g = abs(g);
			
			ll mx = z * z / g;
			mx = max(mx, abs(x * x - (ll)y * y * d) / g);
			mx = max(mx, abs(2ll * x * z) / g);
			return k / mx * 2 + 1;
		}
	}
};