numpy
matrix operations
triangular matrix
python programming
data manipulation

Extract upper or lower triangular part of a numpy matrix

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Extracting the upper or lower triangular part of a matrix is a routine step in numerical work, especially with symmetric matrices such as covariance, correlation, and distance tables. NumPy has direct helpers for this, but the details around diagonal offsets, masks, and memory use still matter. The right method depends on whether you want another matrix, a boolean mask, or just the coordinates of one triangular region.

Start With np.triu and np.tril

NumPy provides two direct functions:

  • 'np.triu for the upper triangle'
  • 'np.tril for the lower triangle'
python
1import numpy as np
2
3A = np.array([
4    [1, 2, 3],
5    [4, 5, 6],
6    [7, 8, 9],
7])
8
9upper = np.triu(A)
10lower = np.tril(A)
11
12print(upper)
13print(lower)

These functions return new arrays where the values outside the chosen triangle are replaced with zero. That is convenient for many matrix algorithms, but it is worth remembering that they are not in-place views into the original array.

Control the Diagonal With k

The k argument shifts which diagonal is included. This is often the difference between the correct result and an off-by-one error.

python
1strict_upper = np.triu(A, k=1)
2strict_lower = np.tril(A, k=-1)
3
4print(strict_upper)
5print(strict_lower)

Use k=0 when the main diagonal should remain, k=1 for a strict upper triangle, and k=-1 for a strict lower triangle. That matters in pairwise matrices where the diagonal may represent self-relations you want to exclude.

Use Masks When You Need Selection Instead of Zeroes

Sometimes a zero-filled matrix is not the output you want. You may need a mask for indexing, filtering, or custom assignment. In that case, build a triangular mask once and reuse it.

python
1mask = np.triu(np.ones_like(A, dtype=bool), k=1)
2values = A[mask]
3
4print(mask)
5print(values)

This is especially useful when converting a symmetric matrix into a one-dimensional feature vector by keeping only one side of the matrix.

Index Helpers Make Extraction Compact

If you want coordinates rather than a dense matrix, use the triangular index helpers.

python
1rows, cols = np.triu_indices_from(A, k=1)
2values = A[rows, cols]
3
4print(rows)
5print(cols)
6print(values)

This approach is often cleaner than generating a zero-filled matrix and then flattening it. It also makes it obvious that you are selecting a subset of positions rather than preserving the full matrix structure.

Think About Shape and Memory

np.triu and np.tril work on rectangular arrays too, but many mathematical uses of triangular extraction assume a square matrix. If your algorithm requires that, validate it explicitly.

python
1def strict_upper_triangle(matrix: np.ndarray) -> np.ndarray:
2    if matrix.ndim != 2 or matrix.shape[0] != matrix.shape[1]:
3        raise ValueError("expected a square matrix")
4    return np.triu(matrix, k=1)

Memory is the other practical concern. For large matrices, building multiple full-size triangular arrays can be wasteful because roughly half the entries are zero. In repeated workflows, masks or index arrays are often more efficient than allocating many dense outputs.

Use the Right Representation for the Next Step

The best triangular extraction method is the one that matches the next operation:

  • Need another matrix for linear algebra: use np.triu or np.tril.
  • Need to filter or assign selectively: use a boolean mask.
  • Need a compact vector of unique pairwise values: use index helpers.

Thinking in terms of the next step keeps the code simple and avoids unnecessary copies.

Common Pitfalls

The biggest mistake is using the wrong k value and accidentally including or excluding the diagonal. Another is assuming np.triu or np.tril modifies the original matrix in place. Teams also recompute masks in inner loops when the matrix shape is fixed, which wastes time and allocations. On large symmetric inputs, choosing a dense triangular copy when an index vector would do can also waste a surprising amount of memory.

Summary

  • Use np.triu and np.tril for the basic upper and lower triangle operations.
  • Adjust diagonal inclusion explicitly with the k parameter.
  • Use masks or index helpers when you need selection instead of zero-filled matrices.
  • Validate square-matrix assumptions if the downstream math requires them.
  • Choose the representation that best fits the next operation, not just the first one that works.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.