# POLYRT - Editorial

Setter: Allen
Tester & Editorialist: Taranpreet Singh

Easy-Medium

# PREREQUISITES

Prefix sums, Maths

# PROBLEM

Let’s define a function P(x, y) = (x-y)^{2}. For given real numbers a and b let’s define a polynomial of x: Q_{ab}(x) = P(P(P(x, a), a-b), x-b).

Now let’s define F(a, b):

• If Q_{ab}(x) has no real roots, F(a, b) = 0.

• Otherwise let r be the largest real root of Q_{ab}(x).

• If r is rational, represent r = \frac{p}{q} where p and q are integers, \mathbf{gcd}(|p|, q) = 1 and q \gt 0 and let F(a, b) = p + q.

• If r is irrational and it can be represented in the form r = \frac{p + \sqrt{z}}{q} where p, q and z are integers, q \gt 0, z \gt 0, choose the representation with the smallest value of q and let F(a, b) = p + z + q.

• It can be proven that under the constraints of this problem one of the options above happens.

You are given an array of N integers c_i. Find \sum_{i=1}^{N} \sum_{j=1}^{N} F(c_i, c_j).

# QUICK EXPLANATION

• The largest real root of given equation, if it exists, is \displaystyle \frac{2*a+1+\sqrt{4*a-4*b+1}}{2} . The roots are real only when a \geq b.
• If 4*a-4*b+1 is not a perfect square, Q_{a,b}(x) = 6*a-4*b+4. Otherwise Q_{a,b}(x) = (2*a+\sqrt{4*a-4*b+1})/2+1
• There are only \sqrt{MX} candidates for b with given a, if 4*a-4*b+1 is a perfect square.
• Use prefix sums to calculate the answer assuming irrational roots, and subtract the appropriate values for these pairs.

# EXPLANATION

### Roots of the given equation

We have P_{a, b}(x) = (((x-a)^2-(a-b))^2-(x-b))^2 = 0, and we need to find its roots. We have
((x-a)^2-(a-b))^2-(x-b) = 0
(x-a)^2-(a-b) = \pm \sqrt{x-b}
\displaystyle (x-a)^2+b = a \pm \sqrt{x-b}

For \displaystyle (x-a)^2+b = a + \sqrt{x-b}

Now, assume two functions f(x) = (x-a)^2+b and g(x) = \sqrt{x-b}+a. Visualizing curves (x, f(x)) and (x, g(x))\$, we can see that these are two parabolas. The real roots of original equation appears only at intersection points of these two parabolas. By some observation, we can see that the largest root appears on the like x = y, so we solve for x = f(x) and x= g(x)

Solving x = f(x), we get roots \displaystyle x = \frac{2*a+1 \pm \sqrt{4*a-4*b+1}}{2}

For \displaystyle (x-a)^2+b = a - \sqrt{x-b}
Similar to above situation, we can find roots at line y = -x, leading to roots \displaystyle \frac{2*a-1 \pm \sqrt{4*a-4*b-3}}{2}.

It is easy to notice that \displaystyle Q_{a,b}(x) = \frac{2*a+1+\sqrt{4*a-4*b+1}}{2} is the largest root.

Now, roots exists only when 4*a-4*b+1 \geq 0 \implies a \geq b. Hence we need to check only pairs (a, b) where a \geq b, since F(a, b) = 0 for a < b

If 4*a-4*b+1 is not a perfect square, we need to represent Q_{a,b}(x) in form \displaystyle \frac{p+\sqrt z}{q}, we can see that p = 2*a+1 is odd and q is 2, implying q is already the smallest value possible. So F(a, b) = p+z+q = (2*a+1)+(4*a-4*b+1)+2 = 6*a-4*b+4

If d^2 = 4*a-4*b+1 is a perfect square, d must be an odd integer, making 2*a+1+d an even value. So, F(a, b) = \frac{(2*a+1+\sqrt{4*a-4*b+1}}{2} +1.

### Coming back to the problem

Now, we need to compute the sum of F(a, b) over all pairs values from the given array C.

### Observation

For a fixed value of a, there are \sqrt a candidate values for b for which there exist a d satisfying 4*a-4*b+1 = d^2.

### Implementation

So, let’s compute S = (6*a-4*b+4) for all ordered pairs (a, b) of values where a \geq b. Now, for each pair (a, b) where 4*a-4*b+1 is a perfect square, subtract 6*a-4*b+4 and add \displaystyle \frac{(2*a+1+\sqrt{4*a-4*b+1}}{2} +1 to S.

Now,

After sorting array C, we can compute the first part in O(N) time. For each a, let’s iterate over odd d such that 4*a-d*d+1 \geq 0. Each value b = (4*a-d*d+1)/4 gives a perfect square and we can adjust their contribution accordingly.

Since the original array, C can contain duplicates, we can build the count array B where B_i denotes the number of occurrences of i. Pair (a, b) can be chosen in B_a * B_b ways.

In case of any doubt, refer to my implementation below.

# TIME COMPLEXITY

The time complexity is O(MX+N*\sqrt{MX}) per test case.

# SOLUTIONS

Setter's Solution
#include <bits/stdc++.h>
#define FR(i, N) for (int i = 0; i < int(N); i++)
#define all(x) begin(x), end(x)

using namespace std;

using ll = long long;

const int MAXN = (int)(2e5) + 5;
vector<ll> dSquare;
ll perfectSquare[5*MAXN];
ll cnt[MAXN];
void solve() {
FR(i, MAXN) cnt[i] = 0;
int N;
cin >> N;
for (int i = 0; i < N; i++) {
int a;
cin >> a;
cnt[a]++;
}
ll res = 0;
ll pref = 0;
ll prefCnt = 0;
FR(a, MAXN) {
if (cnt[a]) {
res += cnt[a]*cnt[a]*(a+2);
res += cnt[a]*(prefCnt*6*a+pref);
for (ll k : dSquare) {
if (a - k < 0) break;
res -= cnt[a]*cnt[a-k]*(6*a-4*(a-k)+4);
res += cnt[a]*cnt[a-k]*((2*a+1+perfectSquare[4*k+1])/2 + 1);
}
prefCnt += cnt[a];
pref += (-4*a+4)*cnt[a];
}
}
cout << res << '\n';
}
int main() {
cin.tie(0);
cin.sync_with_stdio(0);
for (int i = 1; i*i < 5*MAXN; i++) {
perfectSquare[i*i] = i;
}
for (int i = 1; i < MAXN; i++) {
if (perfectSquare[i*4 + 1]) {
dSquare.push_back(i);
}
}
int T;
cin >> T;
while (T-->0) {
solve();
}
return 0;
}

Tester's Solution
import java.util.*;
import java.io.*;
class Polyroot{
//SOLUTION BEGIN
void pre() throws Exception{}
void solve(int TC) throws Exception{
int N = ni(), MX = (int)5e5;
long[] cnt = new long[1+MX];
for(int i = 0; i< N; i++)cnt[ni()]++;

long ans = 0;
long pref = 0, prevCount = 0;
for(int a = 0; a<= MX; a++){
if(cnt[a] == 0)continue;
//Adding sum of (6*a-4*b+4) for all pairs with current a and b<= a
prevCount += cnt[a];
pref += cnt[a]*a;
ans += cnt[a]*(6*a*prevCount - 4*pref + 4*prevCount);

//Iterating over pairs (a, b) such that 4*a-4*b+1 = d^2 is a perfect square
for(int d = 1; 4*a-d*d+1 >= 0; d+= 2){
int b = 4*a-d*d+1;
hold(b%4 == 0);
b /= 4;
//Pair (a, b) such that (4*a-4*b+1) is a perfect square
ans -= cnt[a]*cnt[b]*(6*a-4*b+4);
ans += cnt[a]*cnt[b]*(1+(2*a+1+d)/2);
}
}
pn(ans);
}
//SOLUTION END
void hold(boolean b)throws Exception{if(!b)throw new Exception("Hold right there, Sparky!");}
static boolean multipleTC = true;
void run() throws Exception{
out = new PrintWriter(System.out);
//Solution Credits: Taranpreet Singh
int T = (multipleTC)?ni():1;
pre();for(int t = 1; t<= T; t++)solve(t);
out.flush();
out.close();
}
public static void main(String[] args) throws Exception{
new Polyroot().run();
}
int bit(long n){return (n==0)?0:(1+bit(n&(n-1)));}
void p(Object o){out.print(o);}
void pn(Object o){out.println(o);}
void pni(Object o){out.println(o);out.flush();}
String n()throws Exception{return in.next();}
String nln()throws Exception{return in.nextLine();}
int ni()throws Exception{return Integer.parseInt(in.next());}
long nl()throws Exception{return Long.parseLong(in.next());}
double nd()throws Exception{return Double.parseDouble(in.next());}

StringTokenizer st;
}

}

String next() throws Exception{
while (st == null || !st.hasMoreElements()){
try{
}catch (IOException  e){
throw new Exception(e.toString());
}
}
return st.nextToken();
}

String nextLine() throws Exception{
String str = "";
try{
}catch (IOException e){
throw new Exception(e.toString());
}
return str;
}
}
}


Feel free to share your approach. Suggestions are welcomed as always.

I tried to use Fft, by finding number of ways to get a certain a-b value and sum of all A’s over all such combinations for each a-b value. But I repeatedly got WA with it. Could it be my implementation or is it the problem with FFT? Because FFT could cause precision errors from what I’ve heard.

Doesn’t only depend on a-b, also depends on a

If you make the cases for both rational as well irrational roots , the final answer will depend on value of a-b, sum of values of a over all cases and the number of ways for the same value of a-b.

If it’s rational, I belive the answer is (sqrt(4a-4b+1) + 2a + 3)/2
If it’s irrational, it’ll be (4(a-b)+2a+4)