如何用Numpy匹配整个4x4数组而非单行？Where()与All()使用疑问

Alright, let's solve this problem where you want to match entire 4x4 arrays in your numpy array A, instead of just single rows. Here's what's going on and how to fix it:

The Problem with Your Original Code #

Your line results = np.where((A==U[0]).all(axis=-1)) checks for matches along the last axis (each individual row of the 4x4 arrays). That's why you're getting single-row matches instead of full 4x4 array matches. We need to adjust the axis we're checking to cover the entire 4x4 matrix.

The Fix: Check Across Both Rows and Columns #

To match an entire 4x4 array, we need to verify that all elements in the sub-array match U[0]. We can do this by using all(axis=(1,2)) instead of all(axis=-1)—this checks consistency across both the row (axis=1) and column (axis=2) dimensions of each 4x4 sub-array.

Here's the complete adjusted code:

import numpy as np

# Original setup
A = np.random.randint(5, size=(25, 4, 4))
U = np.unique(A, axis=0)

# Step 1: Find indices of all full 4x4 arrays matching U[0]
full_match_indices = np.where((A == U[0]).all(axis=(1, 2)))[0]

# Step 2: Format results to match your desired output (matrix index + row index)
# Repeat each matching matrix index 4 times (one for each row in the 4x4 array)
row_part = np.repeat(full_match_indices, 4)
# Tile [0,1,2,3] to match the number of full matches
col_part = np.tile(np.arange(4), len(full_match_indices))

# Final results in your requested format
results = (row_part, col_part)

How This Works #

• (A == U[0]).all(axis=(1,2)) produces a 1D boolean array of length 25, where each entry is True if the corresponding 4x4 sub-array in A is identical to U[0].
• full_match_indices gives us the positions of these full matches (e.g., [1, 97] in your example).
• We then expand these indices to match your desired output: each full match gets 4 entries (one for each row in the 4x4 array), paired with row indices 0-3.

Test Example #

Let's manually create a test case to verify:

# Create A where the 0th and 5th 4x4 arrays are identical
A = np.zeros((25, 4, 4), dtype=int)
A[5] = A[0]
U = np.unique(A, axis=0)

# Run our code
full_match_indices = np.where((A == U[0]).all(axis=(1, 2)))[0]
row_part = np.repeat(full_match_indices, 4)
col_part = np.tile(np.arange(4), len(full_match_indices))

print(results)
# Output: (array([0, 0, 0, 0, 5, 5, 5, 5]), array([0, 1, 2, 3, 0, 1, 2, 3]))

This exactly mirrors the structure of your example result, where full matches are represented by 4 repeated indices.

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