想了两个非暴力的解法:哈希表、前缀+二分查找,写法巧妙的话均能通过。
IO 部分都相同:
n = int(input()) a = list(map(int, input().split())) # 左往右数第i堆有多少苹果 m = int(input()) q = list(map(int, input().split())) # 第qi个苹果属于哪一堆
- 哈希表:根据实现方式和语言的不同,可能会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)
- 前缀:对于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)