Index of users distributed across rooms - scala

Question

I am new to stackoverflow so please guide me if I explain my problem inappropriately.
I have a collection of integers of length N (an Array[Int]) where each element at some index i represents number of users present in room Ri. In room Ri, the users are also indexed such that indexing of users in the first room is from 1 to R1, second room contains R1 + 1 to R1 + R2 and so on.
Now the input is the index of user and I need to find out in which room the user is present.
My solution goes like this:
def getRoomNumber(userIndex: Int, userDistribution: Array[Int]): Int = {
val cummulativeRooms: Array[Int] =
rooms.tail.scanLeft(rooms.head)((x, y) => x + y)
val roomIndex = cummulativeRooms.takeWhile(_ < userIndex).length + 1
roomIndex
}
Here, as the user 4 will be present in room 2 (because rooms have user distribution like this: Room1(U1, U2), Room2(U3, U4, U5)).
My code is working fine for this input. But I have some hidden test cases out of which half of them passed. But later half does not and some even throws an exception.
Can anyone please tell me what is the problem in my code. Also do you have any other algorithm to do it which is more efficient than this.
More Explanation -
Lets say I have 10 users - U1, U2, U3, U4, U5 and we separate them into N rooms in a sequence which is defined by the userDistribution array - Array(2, 3). This means that users U1 and U2 will be present in room 1 & users from U3 to U5 will be present in room 2.
Now, if I want to find out where the user U4 is, so the output should be 2. Inputs are the user index i.e. 4 in this case and userDistribution array - Array(2, 3)
EDIT: Code changed to a function. The inputs are the user index we want to find and userDistributions containing the number of users present in each room sequentially.
EDIT: Constraints are (We don't have to check for these contraints in our code) -
Both N and Ui can have values between 1 to 10e5.
Also, the Ui will be less than sum of the the array.

Answer 1

If I understand the problem correctly, you only need to iterate the user's distribution array and keep a sum of how many users you have seen until that sum is greater than or equals to the target user.
You can do that pretty easily using an imperative solution:
// Return Option because we can not guarantee the constraints.
// Also, ArraySeq is just an immutable array.
def imperativeSolution(userDistribution: ArraySeq[PositiveInt])(targetUser: PositiveInt): Option[PositiveInt] = {
var acc = 0
for (i <- userDistribution.indices) {
acc += userDistribution(i)
if (targetUser <= acc) return Some(i + 1)
}
None
}
However, this is quite "ugly" because of the mutability and the use of return.
We may rather use a recursive approach:
// It is better to use a recursive data structure like List for a recursive algorithm.
def recursiveSolution(userDistribution: List[PositiveInt])(targetUser: PositiveInt): Option[PositiveInt] = {
#annotation.tailrec
def loop(remaining: List[PositiveInt], idx: PositiveInt, acc: PositiveInt): Option[PositiveInt] =
remaining match {
case x :: xs =>
val newAcc = acc + x
if (targetUser <= newAcc) Some(idx)
else loop(remaining = xs, idx + 1, newAcc)
case Nil =>
None
}
loop(remaining = userDistribution, idx = 1, acc = 0)
}
This is immutable but it has a lot of boilerplate, which we may reduce and be more expressive using a more functional solution:
def functionalSolution(userDistribution: IterableOnce[PositiveInt])(targetUser: PositiveInt): Option[PositiveInt] =
userDistribution
.iterator
.scanLeft(0)(_ + _)
.zipWithIndex
.collectFirst {
case (acc, idx) if (targetUser <= acc) => idx
}
You can see them running here.