Google Code Jam 2022 Round 1C : B. Squary

Aの方が手間取った。
https://codingcompetitions.withgoogle.com/codejam/round/0000000000877b42/0000000000afdf76

問題

整数列EがSquaryであるとは、Eの各要素の総和の二乗と、Eの各要素の二乗の総和が等しいことを示す。

ここで、数列Eと正整数Kが与えられる。
EにK個以下の整数を加え、EがSquaryにできるか。
可能ならその1例を示せ。

解法

K=1のケースを考える。
仮にE=[A,B,C]とし、ここにDを追加すると考えると、

• Eの各要素の総和の2乗：(A^2+B^2+C^2)+2(AB+BC+CA)+2D(A+B+C)+D^2
• Eの各要素の二乗の総和：A^2+B^2+C^2+D^2

よって両者が等しいためには、AB+BC+CA=-D(A+B+C)である必要がある。
これを応用すると、Eの各要素の総和をS1、異なる2要素の積の総和をS2とし、追加する要素をXとすると、S2=-X*S1であればよいことがわかる。
あとはこれを満たすXを答えればよい。

K=2の場合、追加する要素をX,Yとすると、
2*S2+2(X+Y)S1+(X^2+Y^2)=0
を満たす整数X,Yがあればよい。Xをsum(E)程度の範囲まで総当たりしながら、条件を満たすYを探そう。
Editorial曰く、これは必ず解があるようだ。

int N,K;
ll E[1010];

void solve(int _loop) {
	int f,i,j,k,l,x,y;
	
	cin>>N>>K;
	ll s1=0,s2=0;
	FOR(i,N) {
		cin>>E[i];
		s2+=E[i]*s1;
		s1+=E[i];
	}
	
	if(K==1) {
		if(s1==0) {
			if(s2==0) {
				cout<<"Case #"<<_loop<<": 0"<<endl;
			}
			else {
				cout<<"Case #"<<_loop<<": IMPOSSIBLE"<<endl;
			}
		}
		else if(s2%s1==0) {
			ll D=-s2/s1;
			cout<<"Case #"<<_loop<<": "<<D<<endl;
			
		}
		else {
			cout<<"Case #"<<_loop<<": IMPOSSIBLE"<<endl;
		}
	}
	else {
		for(int d=-2000000;d<=2000000;d++) {
			ll t2=s2+d*s1;
			ll t1=s1+d;
			if(t1==0) {
				if(t2==0) {
					cout<<"Case #"<<_loop<<": 0 "<<d<<endl;
					return;
				}
			}
			else if(t2%t1==0) {
				ll D=-t2/t1;
				cout<<"Case #"<<_loop<<": "<<D<<" "<<d<<endl;
				return;
			}
		}
		cout<<"Case #"<<_loop<<": IMPOSSIBLE"<<endl;
	}
}

まとめ

必ず解があるの、証明せず通してしまった…。