Train-Test Split with a Seeded Shuffle
Google● MediumFULLTIMEONSITE INTERVIEWProblem statement
You have n aligned dataset rows numbered from 0 to n - 1. Split their indices into training and test groups without reusing a row. The returned two-row array contains the training indices first and test indices second.
For this deterministic practice variant, initialize the index array in ascending order. Starting at its last position and moving left, advance a 32-bit unsigned state by state = (1664525 * state + 1013904223) mod 2^32; swap the index at the current position with the index at state mod (currentPosition + 1). Return the first n - testCount shuffled indices as training and the final testCount as test. Preserve their shuffled order within each group.
The indices can then be applied to features and labels together, keeping their rows aligned. This exercise does not request a stratified split or a particular ML library.
Function
splitTrainTestIndices(n: int, testCount: int, seed: int) → int[][]Examples
Example 1
n = 5testCount = 2seed = 0return = [[4,0,1],[2,3]]The seeded swap schedule yields [4,0,1,2,3]; the final two shuffled indices form the test set.
Example 2
n = 4testCount = 1seed = 7return = [[1,3,0],[2]]The shuffled order is [1,3,0,2], so the last index is held out.
Example 3
n = 3testCount = 2seed = 1return = [[2],[1,0]]A two-row test set leaves one training row, with no repeated index.
Constraints
2 <= n <= 10000.1 <= testCount < n.0 <= seed <= 2147483647.- Every input row index must occur exactly once across the two returned groups.