MIX2 - Editorial

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

\sum_{i=1}^N \sum_{j=i+1}^N (A_i\times A_j)

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:

(A_1 + A_2 + \ldots + A_N)^2 = \sum_{i=1}^N A_i^2 + 2\cdot\left(\sum_{i=1}^N\sum_{j=i+1}^N \left(A_i\times A_j\right)\right)

Rearranging this a bit, we see that

\sum_{i=1}^N\sum_{j=i+1}^N \left(A_i\times A_j\right) = \frac{(A_1 + A_2 + \ldots + A_N)^2 - \sum_{i=1}^N A_i^2}{2}

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)