[清华集训 2014] Sum

求 \(\sum\limits_{d=1}^{n} (-1)^{\lfloor d\sqrt{r} \rfloor}\) 的值。

\(T \leq 10^4\) 组数据，\(n \leq 10^9\)，\(r \leq 10^4\)。

首先特判 \(\sqrt{r}\) 为整数的情况，若 \(\sqrt{r}\) 为偶数则答案为 \(n\)，否则根据 \(n\) 的奇偶性，\(n\) 为奇数则答案为 \(-1\)，否则答案为 \(0\)。

考虑等式 \((-1) ^ x = 1 - 2(x \bmod 2) = 1 - 2(x - 2\lfloor \dfrac{x}{2} \rfloor) = 1 - 2x + 4\lfloor \dfrac{x}{2} \rfloor\)，此时原式可以化为 \(n- 2\sum\limits_{d=1}^{n} \lfloor d\sqrt{r} \rfloor + 4\sum\limits_{d=1}^{n} \lfloor \dfrac{d\sqrt{r}}{2} \rfloor\)。令 \(f(a,b,c,n)=\sum\limits_{i=1}^{n} \lfloor \dfrac{a\sqrt{r}+b}{c} \times i \rfloor\)，则原题答案为 \(n-2 \times f(1,0,1,n)+4 \times f(1,0,2,n)\)。

考虑如何求出 \(f(a,b,c,n)\)。我们不妨令 \(t=\dfrac{a\sqrt{r}+b}{c}\)，则 \(f(a,b,c,n)=\sum\limits_{i=1}^{n} \lfloor i \times t \rfloor\)，可以用类欧解决。

当 \(t \geq 1\) 时，记 \(p=\lfloor t \rfloor\)，则有 \(f(a,b,c,n)=\sum\limits_{i=1}^{n} \lfloor i \times (t-p) + i \times p \rfloor=\sum\limits_{i=1}^{n} \lfloor i \times \dfrac{a\sqrt{r}+b-pc}{c} \rfloor + p \times \dfrac{n(n+1)}{2}=f(a,b-pc,c,n)+p \times \dfrac{n(n+1)}{2}\)。

当 \(t < 1\) 时，考虑直线 \(y=tx\) 不会经过任何整点，则会产生贡献的最大点为 \((n,\lfloor nt \rfloor)\)。考虑容斥，可将答案化为 \(n \times \lfloor nt \rfloor - \sum\limits_{i=1}^{\lfloor nt \rfloor} \lfloor \dfrac{i}{t} \rfloor=n \times \lfloor nt \rfloor - \sum\limits_{i=1}^{\lfloor nt \rfloor} \lfloor i \times \dfrac{c}{a\sqrt{r}+b} \rfloor=n \times \lfloor nt \rfloor - \sum\limits_{i=1}^{\lfloor nt \rfloor} \lfloor i \times \dfrac{c \times (a\sqrt{r}-b)}{(a\sqrt{r}+b) \times (a\sqrt{r}-b)} \rfloor=n \times \lfloor nt \rfloor - \sum\limits_{i=1}^{\lfloor nt \rfloor} \lfloor i \times \dfrac{ac\sqrt{r}-bc}{a^2r-b^2} \rfloor=n \times \lfloor nt \rfloor - f(ac,-bc,a^2r-b^2,\lfloor nt \rfloor)\)。

注意精度，注意约分。

#include<iostream>
#include<cstdio>
#include<cmath>
#include<algorithm>
using namespace std;
long double real_sqrt;
long long r;
long long my_gcd(long long num1,long long num2){
	while(num2){
		long long num3=num1%num2;
		num1=num2;
		num2=num3;
	}
	return num1;
}
long long f(long long a,long long b,long long c,long long n){
	if(n==0){
		return 0;
	}
	long long gcd=my_gcd(a,my_gcd(b,c));
	a/=gcd;
	b/=gcd;
	c/=gcd;
	long double t=(a*real_sqrt+b)/c;
	if(t>=1){
		long long p=round(floor(t));
		return f(a,b-p*c,c,n)+p*n*(n+1)/2;
	}
	else{
		long long p=round(floor(n*t));
		return n*p-f(a*c,-b*c,a*a*r-b*b,p);
	}
}
int main(){
	int T;
	scanf("%d",&T);
	while(T--){
		long long n;
		scanf("%lld %lld",&n,&r);
		real_sqrt=sqrtl(r);
		long long maybe_sqrt=round(real_sqrt);
		if(maybe_sqrt*maybe_sqrt==r){
			if(maybe_sqrt%2==0){
				printf("%lld\n",n);
			}
			else{
				if(n%2==1){
					printf("-1\n");
				}
				else{
					printf("0\n");
				}
			}
		}
		else{
			printf("%lld\n",n-2*f(1,0,1,n)+4*f(1,0,2,n));
		}
	}
    return 0;
}