BOJ 25399 라그랑주님 수학에는 뺄셈도 있어요

문제 링크입니다.

문제 요약

정수 $N$을 서로 다른 양의 제곱수들의 합 또는 차로 표현하세요. 이때 사용되는 제곱수의 개수를 최소화해야 합니다.

풀이

$N < 0$이면 마지막에 부호를 뒤집어주면 되니 일반성을 잃지 않고 $N > 0$이라고 합시다. 이를 몇 가지 케이스로 나눌 수 있습니다. 순서대로,

• $N = n^2$이면 답은 자명하게 1개의 제곱수로 표현할 수 있습니다.
• $N$이 홀수라면 인접한 두 제곱수의 차가 $1,3,5,7,\cdots$ 꼴로 홀수들이 순서대로 나오는 것을 이용해 2개의 제곱수로 표현할 수 있습니다.
• $N$이 $4$의 배수라면 $(t+1)^2 - (t-1)^2 = 4t$임을 이용해 2개의 제곱수로 표현할 수 있습니다.
• 남은 경우는 전수조사를 통해 2개의 양의 제곱수로 $N$을 분할할 수 있는지 직접 확인해 봅시다.
• 만약 이것도 불가능하다면, 3개의 양의 제곱수로 $N$을 분할하는 방법을 찾아야 합니다. 이때 위에서 처리한 조건들을 제외하면 남는 것이 $N = 4k+2$일 때뿐이므로, 적당히 큰 홀수 제곱수에서 $N$을 빼 남은 값이 홀수이도록 만들어 다른 두 제곱수로 처리해 줄 수 있습니다.

이상을 잘 구현하면 AC를 받을 수 있습니다.

코드

#include<bits/stdc++.h>
using namespace std;
using ll = long long;

vector<ll> p, m;
bool rev;

int main(){
	cin.tie(0);cout.tie(0);ios::sync_with_stdio(false);
	ll n; cin >> n;
	if(!n) return !(cout << "3\n+ 3\n+ 4\n- 5");
	if(n < 0){
		rev = 1; n *= -1;
	}
	if((ll)sqrt(n)*(ll)sqrt(n) == n) p.push_back((ll)sqrt(n));
	else if(n&1){
		p.push_back((n+1)/2);
		m.push_back((n+1)/2-1);
	}
	else if(!(n&3)){
		p.push_back((n>>2)+1);
		m.push_back((n>>2)-1);
	}
	else{
		bool done = 0;
		for(ll i = 1; i <= 1000000; i++){
			ll tmp = n-i*i;
			if(tmp <= i*i) break;
			if((ll)sqrt(tmp)*(ll)sqrt(tmp) == tmp){
				p.push_back((ll)sqrt(tmp));
				p.push_back(i);
				done = 1;
				break;
			}
		}
		if(!done){
		p.push_back(1000003);
		n = 1000003LL*1000003LL-n;
		m.push_back((n+1)/2);
		p.push_back((n+1)/2-1);
		}
	}
	ll ans = 0;
	if(rev) swap(p, m);
	cout << p.size()+m.size() << '\n';
	for(auto i : p) cout << "+ " << i << '\n';
	for(auto i : m) cout << "- " << i << '\n';
}