Sum of Subarrays#12232

Sum of Subarrays (Editorial)

Written By : MathModel

We are asked to answer queries on an integer sequence $a = (a_{1}, a_{2}, \dots, a_{n})$ a=(a_1,a_2,\dots,a_n). For each query $(l_{i}, r_{i})$ (l_i,r_i) let $L = l_{i}, R = r_{i}$ L=l_i,\ R=r_i. We must compute:

l = L \sum R r = l \sum R j = l \sum r a_{j}

\sum_{l=L}^{R}\ \sum_{r=l}^{R}\ \sum_{j=l}^{r} a_j

i.e. the sum of sums of all non-empty subarrays fully inside $[L, R]$ .

Claim #1 : For a fixed query $[L, R]$ , each element $a_{j}$ a_j contributes independently, by a count equal to the number of subarrays inside $[L, R]$ that contain position $j$ .

Proof : In the expression

l = L \sum R r = l \sum R j = l \sum r a_{j},

\sum_{l=L}^{R}\sum_{r=l}^{R}\sum_{j=l}^{r} a_j,

the inner sum adds $a_{j}$ a_j exactly when $j$ lies inside the chosen subarray $[l, r]$ [l,r]. So we can swap the order of summation: instead of summing over subarrays first, we sum over positions $j$ , and count how many subarrays include $j$ . Thus the total equals

j = L \sum R a_{j} \cdot (number of subarrays [l, r] \subseteq [L, R] with l \leq j \leq r) .

\sum_{j=L}^{R} a_j \cdot (\text{number of subarrays }[l,r]\subseteq[L,R]\text{ with } l\le j\le r).

Now we only need to count those subarrays for each $j$ .

A subarray $[l, r]$ [l,r] within $[L, R]$ contains $j$ iff

L \leq l \leq j and j \leq r \leq R .

L\le l\le j \quad\text{and}\quad j\le r\le R.

So:

$l$ has $(j - L + 1)$ (j-L+1) choices.
$r$ has $(R - j + 1)$ (R-j+1) choices.

Therefore the number of subarrays containing $j$ is

(j - L + 1) (R - j + 1) .

(j-L+1)(R-j+1).

So the answer for query $[L, R]$ is:

Ans (L, R) = j = L \sum R a_{j} (j - L + 1) (R - j + 1) .

\text{Ans}(L,R)=\sum_{j=L}^{R} a_j (j-L+1)(R-j+1).

Claim #2 : The weight $(j - L + 1) (R - j + 1)$ (j-L+1)(R-j+1) can be expanded into a quadratic polynomial in $j$ .

Proof : Rewrite the weight as:

(j - L + 1) (R - j + 1) = (j - (L - 1)) ((R + 1) - j) .

(j-L+1)(R-j+1)=(j-(L-1))((R+1)-j).

Expanding:

(j - (L - 1)) ((R + 1) - j) = j (R + 1 - j) - (L - 1) (R + 1 - j) = j (R + 1) - j^{2} - (L - 1) (R + 1) + (L - 1) j = - j^{2} + (L + R) j + (R + 1) (1 - L) .

\begin{align*} (j-(L-1))((R+1)-j) &= j(R+1-j)-(L-1)(R+1-j)\\ &= j(R+1)-j^2-(L-1)(R+1)+(L-1)j\\ &= -j^2+(L+R)j+(R+1)(1-L). \end{align*}

Thus:

(j - L + 1) (R - j + 1) = - j^{2} + (L + R) j + (R + 1) (1 - L) .

(j-L+1)(R-j+1)= -j^2+(L+R)j+(R+1)(1-L).

Substitute this back:

Ans (L, R) = j = L \sum R a_{j} (- j^{2} + (L + R) j + (R + 1) (1 - L)) = - j = L \sum R a_{j} j^{2} + (L + R) j = L \sum R a_{j} j + (R + 1) (1 - L) j = L \sum R a_{j} .

\begin{align*} \text{Ans}(L,R) &=\sum_{j=L}^R a_j\Big(-j^2+(L+R)j+(R+1)(1-L)\Big)\\[4pt] &= -\sum_{j=L}^R a_j j^2 \ +\ (L+R)\sum_{j=L}^R a_j j \ +\ (R+1)(1-L)\sum_{j=L}^R a_j. \end{align*}

So if we can answer the three segment sums

\sum a_{j}, \sum a_{j} j, \sum a_{j} j^{2}

\sum a_j,\qquad \sum a_j j,\qquad \sum a_j j^2

in $O (1)$ per query, we are done.

Define prefix sums (1-indexed):

$p_{0} [i] = \sum_{k = 1}^{i} a_{k}$ p_0[i]=\sum_{k=1}^i a_k
$p_{1} [i] = \sum_{k = 1}^{i} a_{k} \cdot k$ p_1[i]=\sum_{k=1}^i a_k\cdot k
$p_{2} [i] = \sum_{k = 1}^{i} a_{k} \cdot k^{2}$ p_2[i]=\sum_{k=1}^i a_k\cdot k^2

They are precomputed in $O (n)$ .

For a query $[L, R]$ let:

$s_{0} = p_{0} [R] - p_{0} [L - 1]$ s_0=p_0[R]-p_0[L-1]
$s_{1} = p_{1} [R] - p_{1} [L - 1]$ s_1=p_1[R]-p_1[L-1]
$s_{2} = p_{2} [R] - p_{2} [L - 1]$ s_2=p_2[R]-p_2[L-1]

Then:

Ans (L, R) = - s_{2} + (L + R) s_{1} + (R + 1) (1 - L) s_{0}

\text{Ans}(L,R)= -s_2+(L+R)s_1+(R+1)(1-L)s_0

n,q=map(int,input().split())
a=list(map(int,input().split()))

p=[0]*(n+1)
p2=[0]*(n+1)
p3=[0]*(n+1)

for i in range(n):
    p[i+1]=p[i]+a[i]
    p2[i+1]=p2[i]+a[i]*(i+1)
    p3[i+1]=p3[i]+a[i]*(i+1)*(i+1)

for _ in range(q):
    l,r=map(int,input().split())
    s0=p[r]-p[l-1]
    s1=p2[r]-p2[l-1]
    s2=p3[r]-p3[l-1]
    ans=-s2+(l+r)*s1+(r+1)*(1-l)*s0
    print(ans)

Python

Complexity: $O (n + q)$ \mathcal{O}(n+q)