适用于坐标索引与高效维度求和的稀疏对象选型咨询

针对你需要的高维稀疏标量容器（支持d维索引、默认值0、高效维度求和），提供以下两种可行方案：

方案一：TensorLy 稀疏张量（推荐） #

TensorLy的SparseTensor原生支持任意维度的稀疏数据存储，完全匹配你的需求，且底层做了性能优化，适合大规模稀疏场景。

步骤与代码示例 #

1. 安装依赖

pip install tensorly

2. 实现代码

import tensorly as tl
from tensorly.sparse import SparseTensor
import numpy as np

# 示例数据：合并重复坐标的数值
data = [((1,0,5),6),((2,6,5),100),((5,3,1),1),((2,0,5),4),((2,6,5),100)]
coord_counts = {}
for coords, val in data:
    coord_counts[coords] = coord_counts.get(coords, 0) + val

# 转换为TensorLy要求的格式：坐标为二维数组（每行一个坐标），数值为一维数组
coords = np.array(list(coord_counts.keys()), dtype=int)
values = np.array(list(coord_counts.values()), dtype=int)

# 定义张量形状（示例为3维，每个维度大小10，可扩展到5-100维）
shape = (10, 10, 10)

# 创建稀疏张量
sparse_tensor = SparseTensor((coords.T, values), shape=shape)

# 高效执行维度求和，用法与numpy完全一致
x = sparse_tensor.sum(axis=(1,2))
y = sparse_tensor.sum(axis=(0,))
z = sparse_tensor.sum(axis=(0,2))

# 验证结果与numpy数组一致
np_data = np.zeros(shape)
for coords, val in coord_counts.items():
    np_data[coords] = val
assert np.array_equal(x, np_data.sum(axis=(1,2)))
assert np.array_equal(y, np_data.sum(axis=(0,)))
assert np.array_equal(z, np_data.sum(axis=(0,2)))

方案优势 #

• 支持5-100维的任意高维索引
• 仅存储非零元素，内存占用远低于numpy稠密数组
• 求和等操作接口与numpy对齐，学习成本低
• 底层优化过运算性能，适合高频求和场景（如示例中的循环求和）

方案二：字典自定义实现（轻量无依赖） #

如果不想引入第三方库，可以用字典存储非零坐标的数值，手动实现维度求和逻辑，适合简单场景。

代码示例 #

import numpy as np

# 示例数据：合并重复坐标的数值
data = [((1,0,5),6),((2,6,5),100),((5,3,1),1),((2,0,5),4),((2,6,5),100)]
coord_dict = {}
for coords, val in data:
    coord_dict[coords] = coord_dict.get(coords, 0) + val

# 定义张量形状
shape = (10, 10, 10)

# 自定义维度求和函数
def sum_over_dims(coord_dict, shape, axes):
    result_dict = {}
    # 遍历所有非零坐标，合并求和维度的数值
    for coords, val in coord_dict.items():
        # 保留非求和维度的坐标值
        new_coords = tuple(coord for idx, coord in enumerate(coords) if idx not in axes)
        result_dict[new_coords] = result_dict.get(new_coords, 0) + val
    # 转换为与numpy格式一致的数组
    result_shape = tuple(s for idx, s in enumerate(shape) if idx not in axes)
    result = np.zeros(result_shape, dtype=int)
    for coords, val in result_dict.items():
        result[coords] = val
    return result

# 执行求和操作
x = sum_over_dims(coord_dict, shape, axes=(1,2))
y = sum_over_dims(coord_dict, shape, axes=(0,))
z = sum_over_dims(coord_dict, shape, axes=(0,2))

# 验证结果
np_data = np.zeros(shape)
for coords, val in coord_dict.items():
    np_data[coords] = val
assert np.array_equal(x, np_data.sum(axis=(1,2)))

方案优势 #

• 无额外依赖，代码轻量
• 逻辑简单易懂，可按需修改扩展
• 适合小规模稀疏数据场景

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