如何将二维矩阵的1广播为独热三维矩阵？（Julia/NumPy场景）

Let's break this down into two common scenarios: the standard category-based one-hot encoding, and the specific case shown in your example. Both can be solved cleanly with broadcasting, which works similarly in Julia and NumPy.

Scenario 1: Standard Category-Based One-Hot Encoding #

If your 2D matrix A contains category indices (e.g., 0 or 1 in your example), you want a 3D array B where B[i,j,c] = 1 exactly when A[i,j] equals category c.

Julia Implementation #

First, define your matrix and categories. Then use broadcasting to compare each element of A against every category (reshaped to align with the third dimension):

# Your input matrix
A = [[1 1 1 1 0]; [0 0 0 0 1]]

# Define your categories (0 and 1 in this case)
categories = 0:1

# Reshape categories to 1x1x2 so broadcasting aligns with A's 2D shape
B = A .== reshape(categories, 1, 1, length(categories))

This gives a 2x5x2 array where:

• B[:,:,1] has 1s wherever A is 0
• B[:,:,2] has 1s wherever A is 1

NumPy Equivalent #

The logic is identical—we just use NumPy's syntax for adding a new axis:

import numpy as np
A = np.array([[1,1,1,1,0], [0,0,0,0,1]])
categories = np.arange(2)  # 0 and 1

# Add a new axis to categories to enable broadcasting with A
B = A[..., np.newaxis] == categories

Scenario 2: Matching Your Exact Example #

In your example, B is a 2x5x5 array where each slice k has a single 1 at (i,k) if A[i,k] == 1 (and 0 elsewhere). Here's how to generate this with broadcasting:

Julia Implementation #

First, create a matrix of column indices that matches A's shape. Then broadcast a comparison between these indices and the slice indices (reshaped to fit the third dimension), combined with checking if A has a 1:

A = [[1 1 1 1 0]; [0 0 0 0 1]]

# Create a 2x5 matrix where each column k has value k
cols = repeat(1:size(A,2), outer=(size(A,1),1))

# Reshape slice indices to 1x1x5 and broadcast the combined condition
B = (A .== 1) .& (cols .== reshape(1:size(A,2), 1, 1, size(A,2)))

This produces exactly the B you described:

• Slice 1 has a 1 at (1,1) (since A[1,1] ==1 and column 1 matches slice 1)
• Slice 5 has a 1 at (2,5) (since A[2,5] ==1 and column5 matches slice5)

NumPy Equivalent #

Same logic, using NumPy's tiling and new axis syntax:

import numpy as np
A = np.array([[1,1,1,1,0], [0,0,0,0,1]])

# Create a 2x5 matrix of 1-based column indices
cols = np.tile(np.arange(1, A.shape[1]+1), (A.shape[0], 1))

# Broadcast the combined condition
B = (A ==1)[:, :, np.newaxis] & (cols[:, :, np.newaxis] == np.arange(1, A.shape[1]+1))

Key Broadcasting Idea #

The magic here is reshaping one operand to have singleton dimensions (size 1) in the positions where you want to expand across the original matrix. Julia and NumPy automatically stretch these singleton dimensions to match the other operand's shape, letting you perform element-wise operations across all elements and slices without writing loops.

内容的提问来源于stack exchange，提问作者Frank