Python(或NumPy)中与R中的match函数相当的函数

18

在Python中有什么简单的方法可以实现R中match函数的功能吗?
R中match函数的作用是返回第一个参数在第二个参数中第一次匹配的位置向量。

例如,以下是一个R代码片段。

> a <- c(5,4,3,2,1)
> b <- c(2,3)
> match(a,b)
[1] NA NA  2  1 NA

请用Python将其翻译,我正在寻找一个执行以下操作的函数:

>>> a = [5,4,3,2,1]
>>> b = [2,3]
>>> match(a,b)
[None, None, 2, 1, None]

谢谢!

3个回答

27
>>> a = [5,4,3,2,1]
>>> b = [2,3]
>>> [ b.index(x) if x in b else None for x in a ]
[None, None, 1, 0, None]

如果您确实需要基于“一”的位置而不是基于“零”的位置,请添加1。

>>> [ b.index(x)+1 if x in b else None for x in a ]
[None, None, 2, 1, None]

如果您将频繁重复使用此一行代码,可以使其可重用:

>>> match = lambda a, b: [ b.index(x)+1 if x in b else None for x in a ]
>>> match
<function <lambda> at 0x04E77B70>
>>> match(a, b)
[None, None, 2, 1, None]

很好。这是一个很好的方法来避免当元素不存在时出现list.index()错误。 - hatmatrix

2

一种更快的方法,基于Paulo Scardine的答案(随着数组大小的增加,差异变得更有意义)。 如果您不介意失去这个一行代码:

from typing import Hashable, List


def match_list(a: List[Hashable], b: List[Hashable]) -> List[int]:
    return [b.index(x) if x in b else None for x in a]


def match(a: List[Hashable], b: List[Hashable]) -> List[int]:
    b_dict = {x: i for i, x in enumerate(b)}
    return [b_dict.get(x, None) for x in a]


import random

a = [random.randint(0, 100) for _ in range(10000)]
b = [i for i in range(100) if i % 2 == 0]


%timeit match(a, b)
>>> 580 µs ± 15.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit match_list(a, b)
>>> 6.13 ms ± 146 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

match(a, b) == match_list(a, b)
>>> True

1

可以在Python中完成R的匹配功能,并将匹配的索引作为数据帧索引返回(有助于进一步子集化),如下:

import numpy as np
import pandas as pd
def match(ser1, ser2):
"""
return index of ser2 matching elements of ser1(or return np.nan)
equivalent to match function of R
"""
idx=[ser2.index[ser2==ser1[i]].to_list()[0] if ser1.isin(ser2)[i] == True else np.nan for i in range(len(ser1))]
return (pd.Index(idx))

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接