2020网易：丰收[python][哈希][二分查找]

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

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)

2. 前缀：对于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)