Student Sections · Coding Gym

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 space-separated integers \( n \), \( k \) and \( m \). The second line contains space-separated 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:

1
2
3
4
5
6
7


section_of[1..n]
current_student = 1
for section from 1 to m:
    for i from 1 to a[section]:
        section_of[current_student++] = section

print section_of[k]

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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


current_section = 1
section_of[1..n]
for student from 1 to n:

    # if this section is out of slots, go to the next section.
    if a[current_section] == 0:
        current_section++

    # consume one slot for this student
    a[current_section]--
    section_of[student] = current_section

print section_of[k]

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:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11


current_section = 1
for student from 1 to k:

    # if this section is out of slots, go to the next section.
    if a[current_section] == 0:
        current_section++

    # consume one slot for this student
    a[current_section]--

print 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:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16


int t, n, k, m; cin >> t;
vector<int> sections;
while (t--)
{
    cin >> n >> k >> m;
    sections.resize(m);
    copy_n(istream_iterator<int>(cin), m, begin(sections));
    auto sum = 0;
    auto i = 0;
    while (sum < k) // when sum is >= k, found it!
    {
        sum += sections[i];
        i += 1;
    }           
    cout << i << "\n";
}

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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21


int t, n, k, m; cin >> t;
int value;
while (t--)
{
    cin >> n >> k >> m;              
    auto sum = 0;
    auto i = 0;
    while (sum < k)
    {
        cin >> value;
        sum += value;
        i += 1;
        m--;
    }       
    // because of the structure of the test cases
    while (m-->0) 
    {
        cin >> value;
    }
    cout << i << "\n";
}

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 pre-processing.

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:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11


int t, n, k, m; cin >> t;
vector<int> sections;
while (t--)
{
    cin >> n >> k >> m;
    sections.resize(m);
    copy_n(istream_iterator<int>(cin), m, begin(sections));
    partial_sum(begin(sections), end(sections), begin(sections));
    auto it = lower_bound(begin(sections), end(sections), k);
    cout << (distance(begin(sections), it)+1) << "\n";
}

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:

1
2


def studentSection(k, arr):
    return bisect_left(list(itertools.accumulate(arr)), k) + 1

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