如何以向量化方式高效获取超大规模矩阵的唯一值及其索引

Hey there, let's tackle this massive matrix problem efficiently—no loops allowed, obviously, since that'd be a total nightmare for a (150000, 120000) array! Below are two optimized strategies depending on whether you have enough memory to handle the full matrix at once, or need to work around memory constraints.

Memory-Sufficient Scenario: Full Vectorized Processing #

If your system can handle flattening the entire matrix (note: a 150000x120000 int32 matrix takes ~72GB of RAM), this method uses pure NumPy vectorized operations (no Python loops) to get unique values and their indices quickly.

Step-by-Step Code #

import numpy as np

# Your example matrix (scaled down for demonstration)
matrix_labels = np.array([
    [24, 18, 4, 17, 24, 0, 3, 26],
    [21, 11, 14, 9, 3, 27, 18, 14],
    [25, 26, 27, 16, 26, 27, 21, 26],
    [3, 29, 28, 2, 22, 10, 29, 28],
    [21, 29, 1, 5, 6, 7, 8, 9]
])

# 1. Flatten the matrix to 1D (vectorized operation)
flat_matrix = matrix_labels.flatten()

# 2. Get sorted indices of the flattened matrix
sorted_indices = np.argsort(flat_matrix)
sorted_vals = flat_matrix[sorted_indices]

# 3. Find split points between unique values
_, unique_split_indices = np.unique(sorted_vals, return_index=True)

# 4. Split sorted indices into groups per unique value
indices_per_unique_val = np.split(sorted_indices, unique_split_indices[1:])

# 5. Extract the unique values (from the sorted array at split points)
unique_values = sorted_vals[unique_split_indices]

# Optional: Convert flattened indices back to 2D (row, column) positions
unique_value_positions = {}
for val, flat_idx in zip(unique_values, indices_per_unique_val):
    rows, cols = np.unravel_index(flat_idx, matrix_labels.shape)
    unique_value_positions[val] = list(zip(rows, cols))

# Example output
print("Unique values:", unique_values)
print("Positions of value 3:", unique_value_positions[3])

Why This Works #

All operations here are implemented in NumPy's optimized C backend—way faster than any Python loop. Sorting groups identical values together, then splitting the sorted indices gives us clean, vectorized groups for each unique value.

Memory-Constrained Scenario: Block-Based Vectorized Processing #

For a 150000x120000 matrix, full flattening is often impossible due to RAM limits. This approach splits the matrix into smaller blocks, processes each block vectorially, then aggregates results. The only loops here are over blocks (not individual elements), which is negligible in terms of runtime.

Step-by-Step Code #

import numpy as np

# Assume matrix_labels is your massive (150000, 120000) matrix
# Define block size (adjust based on your available RAM)
block_row_size = 15000  # Split into 10 blocks of 15000 rows each
num_blocks = matrix_labels.shape[0] // block_row_size

# Phase 1: Collect all unique values across blocks, then get global unique values
all_block_unique = []
for block_idx in range(num_blocks):
    # Extract a single block (vectorized slicing)
    block = matrix_labels[block_idx*block_row_size : (block_idx+1)*block_row_size, :]
    # Get unique values for the block
    all_block_unique.extend(np.unique(block))
# Global unique values (deduplicated)
global_unique_values = np.unique(all_block_unique)

# Phase 2: Collect indices for each global unique value across blocks
unique_value_indices = {val: [] for val in global_unique_values}

for block_idx in range(num_blocks):
    block = matrix_labels[block_idx*block_row_size : (block_idx+1)*block_row_size, :]
    flat_block = block.flatten()
    
    # For each unique value, find indices in the block (vectorized)
    for val in global_unique_values:
        block_flat_indices = np.where(flat_block == val)[0]
        if len(block_flat_indices) > 0:
            # Convert block-local indices to global flattened indices
            global_flat_indices = block_idx * block_row_size * matrix_labels.shape[1] + block_flat_indices
            unique_value_indices[val].append(global_flat_indices)

# Merge indices for each value into a single array
for val in unique_value_indices:
    unique_value_indices[val] = np.concatenate(unique_value_indices[val])

# Optional: Convert to 2D positions
for val in unique_value_indices:
    rows, cols = np.unravel_index(unique_value_indices[val], matrix_labels.shape)
    unique_value_indices[val] = (rows, cols)

Key Notes #

• Adjust block_row_size based on your RAM (smaller blocks = lower memory usage, slightly more overhead from block loops).
• This method avoids loading the entire matrix into memory at once, making it feasible for ultra-large datasets.

Bonus: Alternative for Out-of-Core Processing #

If even block processing is too memory-heavy, consider using Dask (a parallel computing library that handles out-of-core arrays). It mimics NumPy's API but processes data in chunks, which is perfect for matrices larger than your RAM.

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