像字典一样索引 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 中为模型提供键以及它们的索引？

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

• 如果 key 参数是一个元组或列表，我们会将其中的每个条目视为一个单独的键，以将其转换为 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

*：除非你对模型切片有感觉？