kmjp's blog

競技プログラミング参加記です

yukicoder : No.1918 Simple Math ?

なるほど…
https://yukicoder.me/problems/no/1918

問題

正整数N,aが与えられる。aは10^6以下、Nは10^12以下である。
 \displaystyle \sum_{i=1}^N \lfloor \sqrt{ai} \rfloorを求めよ。

y=√axというグラフを考える。(1,1)-(N,√aN)の矩形領域に含まれる格子点のうち、y=√axより左側にあるものを除けばよい。
これは数えにくいので、x軸とy軸を入れ替え(1,1)-(√aN,N)においてy=(1/√a)*x^2の下側にある格子点を数えよう。

aが小さいことに着目し、計算回数を削減しよう。
x=(ak-(a-1)),(ak-(a-2)),.....akの時のy=(1/√a)*x^2の下側にある点の個数は、kの二乗和や一乗和の線形和として計算できる。
x≦ak≦√aNであるようなxに対しては、格子点の総和を高速に計算できる。

int T;
ll A,N;
const ll mo=1000000007;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		cin>>A>>N;
		
		ll P=A*A;
		ll Q=-2*A*(A-1)/2;
		ll R=0;
		FOR(i,A) {
			R+=1LL*i*i/A-(1LL*i*i%A==0);
		}
		ll M=N/A;
		__int128 r=sqrt(A*N);
		while(r*r>A*N) r--;
		while((r+1)*(r+1)<=A*N) r++;
		__int128 tot=r*N;
		
		ll v=r/A;
		tot-=(__int128)P*v*(v+1)*(2*v+1)/6;
		tot-=(__int128)Q*v*(v+1)/2;
		tot-=(__int128)R*v;
		for(__int128 b=v*A+1;b<=r;b++) tot-=b*b/A-(b*b%A==0);
		tot%=mo;
		cout<<(ll)tot<<endl;
	}
}

まとめ

x軸とy軸をひっくり返すことに考えが至らず…。