Google Code Jam 2022 round1C Squary

Squary
Given N
integer, is it possible to add at least 1 and at most K
elements such that the sum of the squares of all elements in list equals square of sum of the all elements?
solution
Let $$ a_i $$ or A
be the list of the given n
elements.
Let $$ b_j $$ be the list of the elements to add the `A` list.
$$(\sum_{i}a_i + \sum_{j}b_j)^2 = \sum_{i}(a_i)^2 + \sum_{j}(b_j)^2$$
First, K = 1
$$(\sum_{i}a_i + b_0)^2 = \sum_{i}(a_i)^2 + (b_0)^2$$
$$b_0 = \frac{{\sum_{i}(a_i)^2 – (\sum_{i}a_i )^2} }{ 2 \times \sum_{i}a_i} $$
If $$\sum_{i}a_i = 0$$, and $$ (\sum_{i}a_i )^2 = 0$$, then we can choose any value as an answer.
if $$\sum_{i}a_i = 0$$, and $$ (\sum_{i}a_i )^2 \neq 0$$ it is impossible to get a list.
Second, K ≠ 1
It is always possible to get a desired list with at most 2 elements added.
$$ b_0 = 1 – \sum_{i}a_i $$
$$ -b_1 = a_1 \times a_2 + a_1 \times a_3 + … + a_1 \times b_0 + a_2 \times a_3 + a_2 \times a_4 + … a_n \times b_0$$
#include <bits/stdc++.h> using namespace std; using ll = long long; void solve() { int n, k; cin >> n >> k; ll sum_a = 0, sum_b = 0; ll sum_aa = 0, sum_bb = 0; vector<ll> a(n); for (int i = 0; i < n; ++i) { cin >> a[i]; sum_a += a[i]; sum_aa += (a[i] * a[i]); } if (sum_a == 0 && sum_aa == 0) { cout << 1; return; } else if (sum_a == 0) { cout << "IMPOSSIBLE"; return; } ll tt = (sum_aa - pow(sum_a, 2)) / (2 * sum_a); if (tt * (2 * sum_a) == (sum_aa - pow(sum_a, 2))) { cout << tt; return; } else if (k == 1) { cout << "IMPOSSIBLE"; return; } ll t1 = 1 - sum_a; a.push_back(t1); ll t2 = 0; for (int i = 0; i < a.size(); ++i) { for (int j = i + 1; j < a.size(); ++j) { if (i == j) continue; t2 += a[i] * a[j]; } } cout << t1 << " " << -t2; } int main() { cin.tie(0); ios_base::sync_with_stdio(false); int t; cin >> t; for (int tt = 1; tt <= t; ++tt) { cout << "Case #" << tt << ": "; solve(); cout << endl; } return 0; }