## Search

-  Programming   -  Competitive Programming   -  Google   -  CodeJam   -  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;
}