numpy
random sampling
2D array
python
data manipulation

Numpy Get random set of rows from 2D array

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

Selecting random rows from a 2D NumPy array is a common operation in machine learning, simulation, and exploratory analysis. The safest general pattern is to sample row indices first and then slice the array, because that approach is explicit, reproducible, and easy to extend to labels or other aligned arrays.

Sample Row Indices and Slice the Array

The main question is whether sampling should happen with or without replacement. Without replacement means every selected row is unique. With replacement means the same row can appear more than once, which is useful for bootstrap-style workflows.

python
1import numpy as np
2
3rng = np.random.default_rng(42)
4arr = np.arange(30).reshape(10, 3)
5
6k = 4
7indices = rng.choice(arr.shape[0], size=k, replace=False)
8sample = arr[indices]
9
10print("indices:", indices)
11print(sample)

This chooses k row positions from the first axis and returns those rows. It is simple and version-friendly, which is why it remains a strong default even when newer NumPy features exist.

Sampling with Replacement

If you need repeated rows, set replace=True. That is the right choice for resampling methods such as bagging or bootstrap estimation.

python
1bootstrap_idx = rng.choice(arr.shape[0], size=6, replace=True)
2bootstrap_rows = arr[bootstrap_idx]
3
4print("bootstrap indices:", bootstrap_idx)
5print(bootstrap_rows)

The important point is that duplicates are not an error in this mode. They are a deliberate part of the sampling strategy.

Reproducibility with default_rng

Avoid the old global random state when you care about reproducibility. A local generator object is clearer and easier to test.

python
1import numpy as np
2
3rng1 = np.random.default_rng(123)
4rng2 = np.random.default_rng(123)
5
6idx1 = rng1.choice(10, size=5, replace=False)
7idx2 = rng2.choice(10, size=5, replace=False)
8
9print(np.array_equal(idx1, idx2))

Using a dedicated generator also helps when different parts of a pipeline should have isolated random streams.

Alternative: Shuffle Then Take the First k

Another clean approach for sampling without replacement is to generate a random permutation of row indices and take the first k.

python
1import numpy as np
2
3rng = np.random.default_rng(7)
4arr = np.arange(24).reshape(8, 3)
5
6perm = rng.permutation(arr.shape[0])
7sample = arr[perm[:3]]
8
9print("permutation:", perm)
10print(sample)

This is especially readable when you want a full random ordering and only later decide how many rows to keep.

In machine learning code, the array of features is rarely the only thing being sampled. There are usually labels, sample weights, or metadata arrays that must stay aligned with the chosen rows. The safest habit is to sample indices once and then apply them everywhere.

python
1import numpy as np
2
3X = np.arange(60).reshape(20, 3)
4y = np.arange(20)
5weights = np.linspace(1.0, 2.0, 20)
6
7rng = np.random.default_rng(99)
8idx = rng.choice(X.shape[0], size=5, replace=False)
9
10X_sample = X[idx]
11y_sample = y[idx]
12weight_sample = weights[idx]
13
14print(X_sample.shape, y_sample.shape, weight_sample.shape)

Sampling each array independently is a subtle but serious bug because it breaks correspondence between the rows and their targets.

Validate Inputs in Utility Code

If you turn this into a helper function, add basic validation. That makes failure modes clearer for the next person who uses the function.

python
1def sample_rows(a, k, replace=False, seed=None):
2    if a.ndim != 2:
3        raise ValueError("expected a 2D array")
4    if not replace and k > a.shape[0]:
5        raise ValueError("cannot sample more rows than available without replacement")
6
7    rng = np.random.default_rng(seed)
8    idx = rng.choice(a.shape[0], size=k, replace=replace)
9    return a[idx], idx

Returning the indices as well as the sample is often useful for auditing, debugging, or sampling related arrays later.

Common Pitfalls

Trying to sample more rows than exist while using replace=False raises an error. Decide on replacement policy before choosing sample size.

Using the legacy global random API makes experiments harder to reproduce and reason about. Prefer default_rng.

Sampling features and labels separately breaks row alignment and can silently corrupt a training dataset.

Summary

  • The most reliable NumPy pattern is to sample row indices and then slice the 2D array.
  • Choose replace=False for unique rows and replace=True for bootstrap-style sampling.
  • Use np.random.default_rng when reproducibility matters.
  • Sample indices once and reuse them for all aligned arrays such as labels and weights.

Course illustration
Course illustration

All Rights Reserved.