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;
}