python - 像字典一样索引 NumPy 数组

Question

出于性能原因，我需要使用二维 NumPy 数组，但我还需要能够为每个元素建立索引。索引将是 models1 和 models2，它们是 django.db.models.Model 的子类。

我需要能够获取和设置项目、切片和传递索引列表、过滤器等，就像它是具有整数索引的常规 NumPy 一样。但我还需要能够使用提到的索引来做到这一点。

我应该能够执行的操作示例:

arr[m1]
arr[:,m2]
arr[(0,1,4),:] = N

是否有标准化的方法来索引 NumPy 数组？我还没找到。

我已经尝试过这个:

class IndexedArray:
    def __init__(self, models1, models2):
        self.shape = (len(models1), len(models2))
        self.arr = np.full(self.shape, 1)
        self.models1 = [m1 for m1 in models1]
        self.models2 = [m2 for m2 in models2]
        self.model1_indices = {m1: i for i, m1 in enumerate(self.models1)}
        self.model2_indices = {m2: i for i, m2 in enumerate(self.models2)}

    def __setitem__(self, key, value):
        if isinstance(key, tuple):
            i = self.models1_indices[key[0]]
            j = self.models2_indices[key[1]]
            self.arr[i, j] = value
        else:
            i = self.models1_indices[key]
            self.arr[i] = value

    def __getitem__(self, key):
        if isinstance(key, tuple):
            i = self.models1_indices[key[0]]
            j = self.models2_indices[key[1]]
            return self.arr[i, j]
        i = self.models1_indices[key]
        return self.arr[i]

    def get_index(self, o):
        if isinstance(o, Model1):
            return self.models1_indices[o]
        if isinstance(o, Model2):
            return self.models2_indices[o]
        raise TypeError

但是可用的操作数量是有限的。我只能获取和设置 arr[m1,m2], arr[m1] 但不能获取和设置 arr[:,m2] 因为切片不是支持(并不是说它们有意义，但是我如何批量分配给第二个维度？我应该传递什么作为第一个索引？)，仅此而已。也不支持过滤...

Answer 1

如果我理解得很好，您希望能够在字典 modelX_indices 中给出模型键及其索引？

这就是我解决问题的方法:

• 如果关键参数是元组或列表，我们将其中的每个条目视为一个单独的键，以将其转换为 int。
• 如果键是切片，我们不会碰它*。
• 如果键是 Model，我们会将其转换为 int，如果它已经是 int，则不要修改它。

以下函数可以完成此工作:

def key_to_int(self, key):
    if isinstance(key, (tuple, list)):
        return tuple(key_to_int(k) for k in key)
    if isinstance(key, Model1):
        return self.models1_indices[key]
    if isinstance(key, Model2):
        return self.models2_indices[key]
    return key

*:除非您有使用模型进行切片的感觉？