想了两个非暴力的解法:哈希表、前缀+二分查找,写法巧妙的话均能通过。


IO 部分都相同:

n = int(input())
a = list(map(int, input().split())) # 左往右数第i堆有多少苹果
m = int(input())
q = list(map(int, input().split())) # 第qi个苹果属于哪一堆

  1. 哈希表:根据实现方式和语言的不同,可能会TLE。时间复杂度主要取决于苹果的总量。
hm, j = [0] * sum(a), 0
for i in range(n): # 第i堆
    t = a[i]
    while t: # C语言重写这段代码则不会超时
        hm[j] = i
        j += 1
        t -= 1

for i in q:
    print(hm[i-1]+1)

相应的Python经过修改,优化哈希表的构建速度,则可以AC:

hm, j = [], 0
for i in range(n): # 第i堆
    # hm[j:j+a[i]] = [i] * a[i] # 构造方式很重要,该方法比较慢,只能通过90%的数据
    # j += a[i]
    hm += [i] * a[i] # 而该方法快一些,不会TLE

for i in q:
    print(hm[i-1]+1)

  1. 前缀:对于q数组,若想每次查询时无需重复计算前若干项的a[i]之和,则应当把其保存起来。然后二分查找第一个不小于q[i]的index即可。时间复杂度会在O(m*log(n))左右。
# 该方法速度能优化到之前的四分之一
def lower_bound(array, first, last, value): 
    while first < last:
        mid = first + (last - first) // 2  # 防溢出
        if array[mid] < value: 
            first = mid + 1
        else:
            last = mid
    return first 

beg = [0] # 每一堆的起始位置, len(beg)=n
for i in a[:-1]:
    beg += [beg[-1] + i]

for i in q:
    print(lower_bound(beg, 0, n, i)) # 找到beg中大于i的第一个起始位置的index

对于python,还可以直接调STL里的模块:

  • 使用itertools import accumulate来求beg;
  • 使用bisect来二分查找;
from itertools.accumulate
import bisect

n = int(input())
a = list(map(int, input().split())) # 左往右数第i堆有多少苹果
m = int(input())
q = list(map(int, input().split())) # 第qi个苹果属于哪一堆

beg = list(accumulate(a)) # 每一堆的起始位置 len(beg)=n

for i in q:
    print(bisect.bisect_left(beg, i)+1)