Sigma Queries#13037

Sigma Queries (Editorial)

Problem Author:	Hasan Saleh (MathModel)
Prepared by:	Hasan Saleh (MathModel)
Edited by:	Hasan Saleh (MathModel)

Solution

The $- 1$ case never happens , since $a_{i} \geq 1$ a_i \ge 1 the prefix sums of $b$ are strictly increasing and unbounded.

Let $S = i = 1 \sum n a_{i}$ S=\sum\limits_{i=1}^{n} a_i and $P_{r} = i = 1 \sum r a_{i}$ P_r=\sum\limits_{i=1}^{r} a_i , then writing $k = m \cdot n + r$ k=m \cdot n+r with $0 \leq r \leq n$ 0 \le r \le n gives

i = 1 \sum k b_{i} = m \cdot S + P_{r}

\sum\limits_{i=1}^{k} b_i=m \cdot S+P_r

So we only need to choose the number of complete copies $m$ , then finish inside one more copy. Take

m = ⌊ \frac{x - 1}{S} ⌋

m=\left\lfloor \frac{x-1}{S} \right\rfloor

which is the largest number of copies with $m \cdot S < x$ m \cdot S < x , and nothing smaller works because $m^{'} < m$ m' < m and $r < n$ r<n imply

m^{'} \cdot S + P_{r} < (m^{'} + 1) \cdot S \leq m \cdot S < x

m' \cdot S+P_r < (m'+1) \cdot S \le m \cdot S < x

What's left is $y = x - m \cdot S$ y=x-m \cdot S with $1 \leq y \leq S$ 1 \le y \le S , and since $P$ is strictly increasing we binary search the minimum $r$ with $P_{r} \geq y$ P_r \ge y , which exists because $P_{n} = S \geq y$ P_n=S \ge y. The answer is

k = m \cdot n + r

k=m \cdot n+r

Note $r = n$ r=n is allowed and just means the answer lands on the end of a copy.

Complexity is $O (n + q lo g n)$ O(n+q \log n).

def solve():
    n,q=map(int,input().split())
    a=list(map(int,input().split()))
    x=list(map(int,input().split()))
    pre=[0]*(n+1)
    for i in range(1,n+1):
        pre[i]=pre[i-1]+a[i-1]
    total=pre[n]
    res=[]
    for v in x:
        m=(v-1)//total
        y=v-m*total
        lo,hi=1,n
        while lo<hi:
            mid=(lo+hi)//2
            if pre[mid]>=y:
                hi=mid
            else:
                lo=mid+1
        res.append(m*n+lo)
    print(*res)
for _ in range(int(input())):
    solve()

Python