HW2A - Editorial

PROBLEM LINK:
Practice
Source

Author: SQU_Test

DIFFICULTY:
EASY

PREREQUISITES:
Basic programming, partial sum.

PROBLEM:
It is snacks time for Ahmed. His friend, Ali, prepared him a nice game for lunch. Ali brought Ahmed n ordered piles of snacks such that i-th pile contains ai snacks. He labeled all these snacks with consecutive integers: snacks in first pile are labeled with numbers 1 to a_1, snacks in second pile are labeled with numbers a_1 + 1 to a_1 + a_2 and so on. See the example for a better understanding.

Ahmed cannot eat all the snacks (Ali brought a lot) and he is blind, so Ali tells him the labels of the best tasty snacks. Ali will only give Ahmed a snack if Ahmed says correctly in which pile this snack is contained.

Poor Ahmed asks for your help. For all tasty snacks said by Ali, tell Ahmed the correct answers.

EXPLANATION:
There are two solutions:

  1. We can make partial sums ( sum_i = a_1 + a_2 + ... + a_i) and then make a binary search for each query q_i to find the result j with the properties sum_{j - 1} < q_i and sum_j ≥ q_i. This solution has the complexity O(n + m·log(n))
  2. We can pre-calculate the index of the pile for each snack and then answer for each query in O(1). This solution has the complexity O(n + m)

TIME COMPLEXITY:
Time complexity of the first solution is O(n + m.log(n)) while the second solution has a complexity of O(n+m).

SOLUTIONS:

Setter’s First Solution
from bisect import bisect_left

n = int(input())
a = list(map(int,input().split()))
m = int(input())
q = list(map(int,input().split()))

for i in range(1,n):
    a[i] = a[i] + a[i-1]

for i in range(m):
    pos = bisect_left(a, q[i])
    print(pos+1)
Setter’s Second Solution
n = int(input())
a = list(map(int,input().split()))
m = int(input())
q = list(map(int,input().split()))
l = []

for i in range(len(a)):
    for j in range(a[i]):
      l.append(i+1)

for i in range(m):
 print(l[q[i]-1])