SETSZ - Editorial

editorial
, , , , ,
1

PROBLEM LINK:

Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4

Author: raysh07
Tester: sushil2006
Editorialist: iceknight1093

DIFFICULTY:

Simple

PREREQUISITES:

Sorting, prefix sums, binary search

PROBLEM:

You’re given an array A of length N.
Answer Q queries.
Each query gives you X, and you must report the size of the set

\{A_i + j \mid 1 \le i \le N, 1 \le j \le X\}

EXPLANATION:

First, we need to figure out how to answer a single query X.

Observe that we really just have N intervals, of the form [A_i+1, A_i+X] for each 1 \le i \le N.
What we’re interested in is the number of points contained in at least one of these intervals.

To analyze this a bit further, let’s sort the array A in ascending order first, so that A_i \le A_{i+1} for all i (clearly this doesn’t change the answer.)

Now, consider two adjacent points A_i and A_{i+1}.
Observe that:

  • If X \lt A_{i+1} - A_i, then the points A_i+1, A_i+2, \ldots, A_i+X will all be covered by the interval starting at A_i+1.
    However, the points A_i+X+1, A_i+X+2, \ldots, A_{i+1} will not be covered by any interval!
    This is because intervals of length X starting before A_i cannot reach A_i+X+1, while the next interval starting after A_i is [A_{i+1}+1, A_{i+1}+X] so it will also skip every point in [A_i+X+1, A_{i+1}].
    Thus, in the range [A_i+1, A_{i+1}], exactly X points will be covered.
  • If X \ge A_{i+1} - A_i, then all the points in [A_i+1, A_{i+1}] will be covered instead.

So, if we define D_i = A_{i+1} - A_i, the number of points in the range [A_i+1, A_{i+1}] that are covered by an interval will equal exactly \min(X, D_i).
(For convenience, we also define D_N = \infty, to represent the empty space after A_N. The above result still holds.)

Thus, the answer for the given X simply equals the sum of \min(X, D_i) across all i.

We now need to adapt the above idea to answer queries for different X.

This can be done using a combination of sorting, prefix sums, and binary search.

First, note that the sum of \min(X, D_i) doesn’t change even if we rearrange the array D.
So, for convenience, let’s sort D in ascending order too.
For a query X, let m be the smallest index such that D_m \ge X.
Then,

  • For all indices i\ge m, the value of \min(D_i, X) simply equals X.
    There are (N-i+1) such indices, so we obtain a contribution of (N-i+1)\cdot X to the answer.
    (Note that this is using 1-indexing, so if you’re using 0-indexing make sure to adjust appropriately.)
  • For all indices i \lt m, the value of \min(D_i, X) equals D_i instead.
    So, for this section we want the sum D_1 + D_2 + \ldots + D_{i-1}.
    This is simply a prefix sum of the (sorted) array D.
    We can thus compute this also in \mathcal{O}(1) time after building the appropriate prefix sum array.

Hence, for a query X, we’re able to solve the problem in \mathcal{O}(1) time if we know the index m.
Because we sorted D, finding the index m can be done in \mathcal{O}(\log N) time per query, allowing us to hence answer each query in logarithmic time too.

TIME COMPLEXITY:

\mathcal{O}((N+Q)\log N) per testcase.

CODE:

Editorialist's code (PyPy3) 
n, q = map(int, input().split())
a = list(map(int, input().split()))

a.sort()
d = [a[i+1] - a[i] for i in range(n-1)] + [10**9]
d.sort()
pref = [0]*(n+1)
for i in range(1, n+1):
    pref[i] = pref[i-1] + d[i-1]

import bisect
for _ in range(q):
    x = int(input())
    pos = bisect.bisect_left(d, x)
    print((n-pos)*x + pref[pos])