校招真题练习030 丰收(网易)

丰收

题目描述
又到了丰收的季节，恰逢小易去牛牛的果园里游玩。
牛牛常说他对整个果园的每个地方都了如指掌，小易不太相信，所以他想考考牛牛。
在果园里有N堆苹果，每堆苹果的数量为ai，小易希望知道从左往右数第x个苹果是属于哪一堆的。
牛牛觉得这个问题太简单，所以希望你来替他回答。

输入描述:
第一行一个数n(1 <= n <= 105)。
第二行n个数ai(1 <= ai <= 1000)，表示从左往右数第i堆有多少苹果
第三行一个数m(1 <= m <= 105)，表示有m次询问。
第四行m个数qi，表示小易希望知道第qi个苹果属于哪一堆。

输出描述:
m行，第i行输出第qi个苹果属于哪一堆。

def main():
    N = int(input())
    apples = list(map(int,input().split()))
    # N = 5
    # apples = [2,7,3,4,9]

    apples.insert(0,0)
    for i in range(1,N+1):
        apples[i] += apples[i-1]
    
    M = int(input())
    questions = list(map(int,input().split()))
    # print(apples)
    # M = 3
    # questions = [1,25,11]
    for i in range(M):
        cur = questions[i]
        l,r = 0,N
        while l <= r:
            m = l + (r - l) // 2
            if cur == apples[m]:
                print(m)
                break
            elif cur < apples[m]:
                if m - 1 >= 0 and cur > apples[m-1]:
                    print(m)
                    break
                else:
                    r = m - 1
            elif cur > apples[m]:
                if m + 1 <= N and cur < apples[m+1]:
                    print(m+1)
                    break
                else:
                    l = m + 1
            else:
                print('error')

if __name__ == '__main__':
    main()

算法思路：二分查找(变形)

寻找有序数组中右边界值。可以使用：itertools.accumulate()和bisect.bisect_left()方法简化代码。