查找最小数字的索引

Question

我有以下（采样的）字典 A，它最初有超过 17,000 个键，每个数组的长度刚刚超过 600,000（对所有人来说都是一样的）。我试图为 600,000 个输入中的每一个找到数组中最小数字的键。例如，在下面的字典中，我想得到 i = 3093094 for j = 0 因为 45.16672136 是所有数组的第一个索引中最小的。同样，对于 j = 1，i = 1157086 因为 1.53174068 是最小的。

A = {3093094: array([45.16672136,  1.68053313, 13.78822307, ..., 36.18798239,
        36.09565274, 35.85261821]),
 1156659: array([45.46286695,  1.69632425, 13.81351489, ..., 36.54544469,
        36.45329774, 36.20969689]),
 1156667: array([45.43970605,  1.69026244, 13.81365067, ..., 36.51934187,
        36.42716964, 36.18364528]),
 1156792: array([45.29956347,  1.57736575, 13.90834355, ..., 36.43079348,
        36.33804273, 36.09623309]),
 1157086: array([45.38149498,  1.53174068, 13.98398836, ..., 36.57985343,
        36.48684657, 36.2457831 ]),
 1430072: array([45.46114909,  1.58096885, 13.95459557, ..., 36.64775128,
        36.55496457, 36.31324461]),
 1668445: array([45.44073352,  1.5941793 , 13.92953699, ..., 36.60630965,
        36.51361336, 36.27162926]),
 3055958: array([45.45006118,  1.57686417, 13.95499241, ..., 36.63558996,
        36.54278917, 36.30111176]),
 1078241: array([45.56175847,  1.77256163, 13.75586274, ..., 36.61441986,
        36.52264105, 36.27795081])}

我有以下多处理解决方案方法，但由于处理时间太长，正在寻找更有效的方法。

import numpy as np
import os
from multiprocessing import Pool


C = range(len(A[3093094]))

def closest(All_inputs):
    (A,j) = All_inputs
    B = list(A.keys())
    my_list = [A[i][j] for i in B]
    return(B[np.argmin(np.array(my_list))])

with Pool(processes=os.cpu_count()) as pool:
    results = pool.map(closest, [(A,j) for j in C])

一个挑战是在多处理中复制 A，因为它的大小很大。你有什么 Pythonic 方法可以快速完成这个看似微不足道的计算吗？

Answer 1

如果你的内存足够大。也许你可以试试这个，使用熊猫。如果仍然很慢，请尝试使用 dask。这两个示例都在下面列出。

import numpy as np
import pandas as pd
import dask.dataframe as dd
from tqdm import tqdm
test_data = {}
for i in tqdm(range(2000)):
    test_data[i] = np.random.randint(0, 10000, 600000)

# test one
# print(test_data)
now = time.time()
df = pd.DataFrame(test_data)
min_idx = df.idxmin(axis=1)
result_one = dict(zip(range(2000), min_idx.tolist()))
# print(result_one)
print(time.time() - now)

# test two
now = time.time()
df = pd.DataFrame(test_data)
ddf = dd.from_pandas(df, npartitions=multiprocessing.cpu_count())
min_idx = ddf.idxmin(axis=1).compute(scheduler="processes")
result_two = dict(zip(range(2000), min_idx.tolist()))
# print(result_two)
print(time.time() - now)

Answer 2

这似乎有效，并且应该比将每一列转换为具有非 Python 列表理解的 Python 列表然后返回 NumPy 数组更快：

K = np.array(list(A))
V = np.array(list(A.values()))
print(K[V.argmin(axis=0)])

示例数据的输出（删除了...）：

[3093094 1157086 1078241 3093094 3093094 3093094]

Try it online!

Answer 3

我在一台 12 核和 16G RAM 的机器上尝试了以下操作：

from multiprocessing import Pool, cpu_count
from time import perf_counter

def closest(values):
    return np.argmin(np.array(values))

if __name__ == "__main__":
    # Build A inside __main__ (otherwise each process builds it again)
    num_keys = 10_000
    arr_len = 100_000
    rng = np.random.default_rng()
    A = {
        key: rng.integers(0, 1000, arr_len)
        for key in range(1000, 1000 + num_keys)
    }

    # Multiprocessing
    start = perf_counter()
    with Pool(processes=cpu_count()) as p:
        indices = p.imap(closest, zip(*A.values()), chunksize=1000)
        keys = tuple(A.keys())
        results = [keys[i] for i in indices]
    end = perf_counter()
    print(f"Duration (np.argmin mp): {end - start:.2f}")

    # np.argmin directly
    start = perf_counter()
    arr = np.array([*A.values()])
    keys = tuple(A.keys())
    results = [keys[i] for i in np.argmin(arr, axis=0)]
    end = perf_counter()
    print(f"Duration (np.argmin direct): {end - start:.2f}")

持续时间结果：

Duration (np.argmin mp): 1258.07
Duration (np.argmin direct): 563.84

---

小样本的结果（num_keys = 4、arr_len = 8）：

A =
{1000: [879, 130, 114, 973, 691, 394, 122, 215],
 1001: [221, 482, 510, 319, 454, 585, 767, 138],
 1002: [982, 526, 971, 168, 185, 477, 838, 37],
 1003: [675, 293, 769, 878, 611, 695, 237, 129]}
results = [1001, 1000, 1000, 1002, 1002, 1000, 1000, 1002]
results = [1001, 1000, 1000, 1002, 1002, 1000, 1000, 1002]