PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author & Editorialist: iceknight1093
DIFFICULTY:
TBD
PREREQUISITES:
None
PROBLEM:
You’re given the quantities of N different drinks Chef has with him.
Find the number of ways of mixing two drinks.
EXPLANATION:
As mentioned in the statement, mixing two drinks is done by:
- Choosing 1 \leq i \lt j \leq N
- Choosing 1 \leq x \leq A_i and 1 \leq y \leq A_j
Suppose we fix i and j.
How many choices do we have for x and y?
Well, x can be anything from 1 to A_i, for A_i choices in total.
Once x is fixed, we can see that similarly, y has A_j choices.
This gives us a total of A_i\times A_j different mixes that can be made using these two drinks.
That means, the answer is simply
The constraints are low, so this can be directly computed in \mathcal{O}(N^2) to get AC.
However, it’s also possible to solve this task in \mathcal{O}(N) time, using a bit of math.
Recall the identity for the square of a summation of several variables:
Rearranging this a bit, we see that
The RHS is easy to compute in \mathcal{O}(N) time since you just want the sum of squares and the square of the sum of the elements; so the required value can be found in \mathcal{O}(N) time.
TIME COMPLEXITY
\mathcal{O}(N^2) or \mathcal{O}(N) per testcase.
CODE:
Author's code (Python, quadratic)
for _ in range(int(input())):
n = int(input())
a = list(map(int, input().split()))
ans = 0
for i in range(n):
for j in range(i+1, n):
ans += a[i]*a[j]
print(ans)
Author's code (Python, linear)
for _ in range(int(input())):
n = int(input())
a = list(map(int, input().split()))
ans = 0
square_of_sum = sum(a)**2
sum_of_squares = sum(x**2 for x in a)
print((square_of_sum - sum_of_squares)//2)