See the original problem on HackerRank.
The school term starts soon, and the students need to be sorted into their respective sections. There are \( n \) students numbered \( 1 \) to \( n \), and \( m \) sections numbered \( 1 \) to \( m \).
Section needs to have exactly \( a_i \) students. To simplify the process, the school has decided to assign students to sections in increasing student number, which means the first \( a_1 \) students will be assigned section \( 1 \), the next \( a_2 \) students will be assigned section \( 2 \), and so on.
Your student number is \( k \). Which section will you be assigned to?
Input Format
The first line of input contains \( t \), the number of scenarios.
The description of \( t \) scenarios follow, each described in two lines. The first line contains three spaceseparated integers \( n \), \( k \) and \( m \). The second line contains spaceseparated integers \( a_1, a_2, …, a_m \).
Constraints
 \( 1 \leq t \leq 250 \)
 \( 1 \leq k \leq n \leq 500 \)
 \( 1 \leq m \leq 500 \)
 \( 1 \leq a_i \leq n \)
 \( a_i + a_2 + … + a_m = n \)
Output
For each scenario, print a single line containing a single integer denoting the section in which student number \(k\) will be assigned to.
Solutions
We call this kind of problems manifold because there are several different/competing solutions.
Naive solutions (courtesy of HackerRank)
First of all, we could simply count the students and assing them to each section. We do a sort of “simulation” of assigning students to sections.
Such a solution scales linearly with the number of students:


Another approach consists in decrementing by one the section anytime a student is assigned to it:


This solution uses some extra space (section_of
), though.
To avoid using extra space, we can iterate from 1
to k
and only keep track of current_section
. The answer is current_section
:


This latter solution is ((O(k))).
Efficient solutions
A more efficient solution scales with the number of section m
.
A first approach consists in accumulating the number of section capacities until we reach a sum that is greater or equal to k
:


This solution requires to read the entire input first (offline). We can make it online:


Clearly, even though we find the answer, we need to consume the entire input afterwards just because of the structure of the test cases (we cannot “skip” the rest of the input automatically).
Variants
A more general problem asks to answer efficiently to more than one query, not just for one student. As you can see, the naive solutions would not scale well and the last one neither.
Another approach consists in doing some preprocessing.
Running a prefix sum on the section array has the effect of accumulating the number of students along the way.
Then, to find the student k
, we could just use a lower bound  that is the first number not less than k
:


Consider the example 11 24 420 6 9
with k=143
.
First of all, we do the prefix sum:
11 35 455 461 470
This means: students from 0
to 11
are in the first section. Students from 12
to 35
in the second, etc. We are just counting (it’s not a coincidence that counting sort has a prefix sum step).
To find 143
efficiently, we can just do a lower bound.
Clearly, this solution does not work offline. It’s a tradeoff.
In Python we have the bisect_left function in library bisect that works like lower_bound as shown in the Roberto Peruzzo’s solution:


A harder variant consists in handling changes/updates on the section array. This can be solved with a segment tree or similar data structures.