Train/Test Split

Implement a train/test split from scratch with NumPy, the way sklearn.model_selection.train_test_split does but without calling it. Interviewers use this to check that you can shuffle reproducibly and keep features and labels aligned.

Task

Complete train_test_split(X, y, test_size=0.2, seed=0) so it returns (X_train, X_test, y_train, y_test).

• X has shape (n, d) and y has shape (n,); row i of X corresponds to y[i].
• Build an index array 0..n-1 and shuffle it with a numpy Generator seeded by seed (np.random.default_rng(seed)).
• The test set is the first int(n * test_size) shuffled indices; the train set is the remainder.
• Index X and y with those index arrays so each pair stays aligned.

Do not call sklearn or any built-in split helper.

Example

X = np.arange(10).reshape(5, 2)   # rows 0..4
y = np.array([0, 1, 2, 3, 4])
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.4, seed=0)

# default_rng(0) shuffles [0,1,2,3,4] -> [2,4,3,0,1]
# int(5 * 0.4) = 2 test rows -> indices [2, 4]; train indices -> [3, 0, 1]
y_test    # -> array([2, 4])
y_train   # -> array([3, 0, 1])
X_test    # -> array([[4, 5],
          #           [8, 9]])

Constraints

• 1 <= n <= 10^6; use vectorized NumPy indexing, not a Python loop.
• 0 <= test_size <= 1.
• The same seed must always produce the same split.

Approach

Shuffle a permutation of the row indices with a seeded default_rng, then slice it into a test block and a train block. Applying the same index arrays to both X and y keeps every feature row paired with its label, and the seed makes the split reproducible.

Solution

import numpy as np

def train_test_split(X, y, test_size=0.2, seed=0):
    n = len(X)
    indices = np.arange(n)
    rng = np.random.default_rng(seed)
    rng.shuffle(indices)                 # in-place seeded permutation

    n_test = int(n * test_size)
    test_idx = indices[:n_test]
    train_idx = indices[n_test:]

    return X[train_idx], X[test_idx], y[train_idx], y[test_idx]

Walkthrough

On the example with n = 5, test_size = 0.4, seed = 0:

• indices = [0, 1, 2, 3, 4].
• default_rng(0).shuffle(indices) gives [2, 4, 3, 0, 1].
• n_test = int(5 * 0.4) = 2, so test_idx = [2, 4] and train_idx = [3, 0, 1].
• X[test_idx] picks rows 2 and 4 ([[4,5],[8,9]]), y[test_idx] = [2, 4].
• X[train_idx] picks rows 3, 0, 1, y[train_idx] = [3, 0, 1].

Features and labels are indexed by the same arrays, so the pairing is preserved.

Complexity & notes

• Time O(n), space O(n) — the shuffle is O(n) and fancy indexing copies out the selected rows.
• Shuffle indices, never X and y separately: two independent shuffles would break the row-to-label alignment.
• int(n * test_size) truncates toward zero, matching the common convention; with test_size=0.4 and n=5 that is 2 test rows.
• np.random.default_rng(seed) is the modern Generator API and gives a reproducible split; passing the same seed anywhere reproduces the exact partition.
• To make the split stratified you would shuffle within each class separately and concatenate, but plain random splitting is enough here.