加快在第二个数组中搜索数组元素

Question

我有一个非常简单的操作，涉及两个不太大的数组：

1. 对于第一个（更大的）数组中的每个元素，位于位置i
2. 查找它是否存在于第二个（较小的）数组中
3. 如果是，则在第二个数组中找到它的索引：j
4. 将取自第三个数组（与第一个数组长度相同）的浮点数存储在位置i，第四个数组的位置j（与第二个数组长度相同）

下面的 for 块可以工作，但对于不太大的数组 (>10000) 会变得非常慢。

这个实现可以更快吗？

import numpy as np
import random

##############################################
# Generate some random data.
#'Nb' is always smaller then 'Na
Na, Nb = 50000, 40000

# List of IDs (could be any string, I use integers here for simplicity)
ids_a = random.sample(range(1, Na * 10), Na)
ids_a = [str(_) for _ in ids_a]
random.shuffle(ids_a)
# Some floats associated to these IDs
vals_in_a = np.random.uniform(0., 1., Na)

# Smaller list of repeated IDs from 'ids_a'
ids_b = random.sample(ids_a, Nb)
# Array to be filled
vals_in_b = np.zeros(Nb)
##############################################

# This block needs to be *a lot* more efficient
#
# For each string in 'ids_a'
for i, id_a in enumerate(ids_a):
    # if it exists in 'ids_b'
    if id_a in ids_b:
        # find where in 'ids_b' this element is located
        j = ids_b.index(id_a)
        # store in that position the value taken from 'ids_a'
        vals_in_b[j] = vals_in_a[i]

Answer 1

为了捍卫我的方法，这里是权威的实现：

import itertools as it

def pp():
    la,lb = len(ids_a),len(ids_b)
    ids = np.fromiter(it.chain(ids_a,ids_b),'<S6',la+lb)
    unq,inv = np.unique(ids,return_inverse=True)
    vals = np.empty(la,vals_in_a.dtype)
    vals[inv[:la]] = vals_in_a
    return vals[inv[la:]]

(juanpa()==pp()).all()
# True

timeit(juanpa,number=100)
# 3.1373191522434354
timeit(pp,number=100)
# 2.5256317732855678

也就是说，@juanpa.arrivillaga 的建议还可以更好地实施：

import operator as op

def ja():
    return op.itemgetter(*ids_b)(dict(zip(ids_a,vals_in_a)))

(ja()==pp()).all()
# True
timeit(ja,number=100)
# 2.015202699229121

Answer 2

我尝试了 juanpa.arrivillaga 和 Paul Panzer 的方法。第一个是迄今为止最快的。也是最简单的。第二个比我原来的方法快，但比第一个慢得多。它还有一个缺点，即vals[inv_a] = vals_in_a 这一行将浮点数存储到U5 数组中，从而将它们转换为字符串。它可以在最后转换回浮点数，但我会丢失数字（当然，除非我遗漏了一些明显的东西。

以下是实现：

def juanpa():
    dict_ids_b = {_: i for i, _ in enumerate(ids_b)}
    for i, id_a in enumerate(ids_a):
        try:
            vals_in_b[dict_ids_b[id_a]] = vals_in_a[i]
        except KeyError:
            pass

    return vals_in_b


def Paul():
    # 1) concatenate ids_a and ids_b
    ids_ab = ids_a + ids_b
    # 2) apply np.unique with keyword return_inverse=True
    vals, idxs = np.unique(ids_ab, return_inverse=True)
    # 3) split the inverse into inv_a and inv_b
    inv_a, inv_b = idxs[:len(ids_a)], idxs[len(ids_a):]
    # 4) map the values to match the order of uniques: vals[inv_a] = vals_in_a
    vals[inv_a] = vals_in_a
    # 5) use inv_b to pick the correct values: result = vals[inv_b]
    vals_in_b = vals[inv_b].astype(float)

    return vals_in_b