收集满足给定条件的numpy数组元素的Pythonic方法

Question

所以我正在处理两个时间序列之间的离散相关函数，数据为 (xi, ti) 和 (xj,tj)，i 和 j = 1,2,3... 计算每个时间序列的滞后时间(i,j) 点对。然后将这些滞后存储在一个 numpy 数组 dst[i,j] 中，其中每个 (i,j) 元素表示该对的滞后时间。

我现在想收集大于某个值的前 n 个滞后及其 (i, j) 索引，但我希望它们是独立的对，这样没有两对具有相同的 i 或相同的 j 项（所以 (1, 2) 和 (3,2) 不起作用）。

举个简单的例子，假设我有：

dst = np.array([[0.2, 0.5, 0.9, 1.0],
                [2.0, 3.0, 4.0, 5.0],
                [7.0, 8.0, 12.0,13.0]])

我想要前两对滞后大于 3。我首先创建了一个 {(i, j) : lag} 形式的字典，其中包含所有滞后 > 3 的元素，然后按滞后排序价值。

idxi, idxj = np.where(dst>3)
mydict = {}
for i, j in zip(idxi, idxj):
    mydict[(i,j)] = dst[i,j]
mydict = {k: v for k, v in sorted(mydict.items(), key=lambda item: item[1])}

#so now mydict = {(1, 2): 4.0, (1, 3): 5.0, (2, 0): 7.0, (2, 1): 8.0, (2, 2): 12.0, (2, 3): 13.0}

所以前两个独立项将是 (1,2) 和 (2,0)。但我不确定获得前两对的最佳方法，同时还要确保没有两对具有相同的 i 和 j 项。我敢肯定我可以想到一种复杂的方法来做到这一点，但我正在寻找一种更 Pythonic 和快速的方法。我对操作 numpy 数组有点陌生，想知道实现目标的最佳方法。那么如何在这里获得前两个独立的对，有没有办法在不创建排序字典的情况下完成整个过程？

Answer 1

我不确定这算不算太复杂，但我认为它至少有效。

dst = [
    [0.2, 0.5, 0.9, 1.0],
    [2.0, 3.0, 4.0, 5.0],
    [7.0, 8.0, 12.0, 13.0]
]


def find_minimum_value(row):
    for column_index, value in enumerate(row):
        if value > 3:
            return column_index, value


answer = {}
for row_index, row in enumerate(dst):
    values = find_minimum_value(row)
    if values:
        column_index, value = values
        answer[(row_index, column_index)] = value

print(answer)

Answer 2

一种方法是使用掩码数组来避免计算相同的索引：

import numpy as np

dst = np.array([[0.2, 0.5, 0.9, 1.0],
                [2.0, 3.0, 4.0, 5.0],
                [7.0, 8.0, 12.0, 13.0]])


def find_maximums(initial, k=2):
    for _ in range(k):
        # find the minimum index and transform to multi-dimensional index
        arg_min = np.unravel_index(initial.argmin(), initial.shape)
        # mask the whole row and the whole column to avoid same indexes
        initial.mask[arg_min[0], :] = initial.mask[:, arg_min[1]] = True
        yield arg_min


res = list(find_maximums(np.ma.masked_less_equal(dst, 3)))
print(res)

输出

[(1, 2), (2, 0)]