PROBLEM LINK:
Contest Division 1
Contest Division 2
Contest Division 3
Practice
Setter: Allen
Tester & Editorialist: Taranpreet Singh
DIFFICULTY
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;
FastReader in;PrintWriter out;
void run() throws Exception{
in = new FastReader();
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());}
class FastReader{
BufferedReader br;
StringTokenizer st;
public FastReader(){
br = new BufferedReader(new InputStreamReader(System.in));
}
public FastReader(String s) throws Exception{
br = new BufferedReader(new FileReader(s));
}
String next() throws Exception{
while (st == null || !st.hasMoreElements()){
try{
st = new StringTokenizer(br.readLine());
}catch (IOException e){
throw new Exception(e.toString());
}
}
return st.nextToken();
}
String nextLine() throws Exception{
String str = "";
try{
str = br.readLine();
}catch (IOException e){
throw new Exception(e.toString());
}
return str;
}
}
}
Feel free to share your approach. Suggestions are welcomed as always. 