PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: rajat397
Tester & Editorialist: iceknight1093
DIFFICULTY:
1526
PREREQUISITES:
None
PROBLEM:
You’re given an array B of length N consisting of 1's and -1's.
Consider a new array C, such that C_i = 2^{i-1}\cdot B_i for each 1\leq i \leq N.
Suppose P subarrays of C have a strictly positive sum and Q of them have a strictly negative sum.
Find |P - Q|.
EXPLANATION:
Consider a subarray [L, R]. A little observation should tell you that:
- If B_R = 1, then C_L + C_{L+1} + \ldots + C_R \gt 0
- If B_R = -1, then C_L + C_{L+1} + \ldots + C_R \lt 0
Proof
Suppose B_R = 1, so C_R = 2^{R-1}
Recall that the powers of two satisfy 2^0 + 2^1 + \ldots + 2^{R-2} \lt 2^{R-1}, which (by removing a few terms) also means 2^{L-1} + \ldots + 2^{R-2} \lt 2^{R-1}.
This can be rearranged to 2^{R-1} - 2^{R-2} - 2^{R-3} - \ldots - 2^{L-1} \gt 0, and the LHS of this equation is the minimum possible sum of C_L + C_{L+1} + \ldots + C_R (attained when B_L = B_{L+1} = \ldots = B_{R-1} = -1)
The minimum possible sum is positive, so clearly any sum is also going to be positive.
A similar proof holds for the B_R = -1 case.
This makes the problem quite simple:
- If B_i = 1, then all the subarrays ending at i have positive sum. There are exactly i such subarrays (in 1-indexing).
- Similarly, if B_i = -1, then all i subarrays ending at this index have a negative sum.
This allows us to compute P and Q both in \mathcal{O}(N) time (if B_i = 1, add i to P; otherwise add i to Q), after which |P - Q| is easily found.
You may also note that, if we let d = P - Q, then each index i adds exactly B_i \cdot i to d.
So, a simple way to write down the answer is \displaystyle\left|\sum_{i=1}^N i\cdot B_i\right|.
TIME COMPLEXITY
\mathcal{O}(N) per test case.
CODE:
Editorialist's code (Python)
for _ in range(int(input())):
n = int(input())
a = list(map(int, input().split()))
ans = 0
for i in range(n):
ans += a[i] * (i+1)
print(abs(ans))