Train Test Split

Split a dataset into a training set and a test set.

Shuffle the row indices, hold out a fraction for testing, and keep the rest for training. X and y have to stay aligned: row i of X must keep its label y[i] through the shuffle.

Use this exact convention so the split is reproducible:

indices    = np.random.default_rng(seed).permutation(n)
n_test     = int(n * test_size)      # truncate, do not round
test_idx   = indices[:n_test]
train_idx  = indices[n_test:]

Return (X_train, X_test, y_train, y_test).

Shuffling matters. Datasets often arrive sorted by label, and slicing off the last 20% without shuffling can hand you a test set that is entirely one class.

Function signature

def train_test_split(X, y, test_size, seed):

Example

Input

X = np.arange(10).reshape(10, 1)
y = np.arange(10)
X_train, X_test, y_train, y_test = train_test_split(X, y, 0.2, seed=0)
y_train, y_test

Output

(array([2, 7, 3, 5, 9, 0, 8, 1]), array([4, 6]))

Reasoning

int(10 * 0.2) = 2, so the first 2 shuffled indices are the test set and the other 8 are the training set. y uses the same indices as X, so every row keeps its label.

Constraints and notes

• Use np.random.default_rng(seed).permutation(n).
• Truncate n * test_size. Do not round.
• Return (X_train, X_test, y_train, y_test).