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.