为numpy数组提供多个相等条件

时间:2017-02-17 00:14:08

标签: python arrays numpy matrix conditional-statements

假设我有一个大的2D numpy数组(称之为 A ),其中包含0到9之间的整数。

我正在尝试编写一个函数,该函数返回与具有以下特征的 A 形状相同的二进制numpy数组(称为 B )。

如果 A 中的相应元素出现在给定列表 L 中,则 B 中的条目为1;否则,条目为零。

下面显示了可能不是最有效的代码。

A = np.random.randint(0, 10, (5, 5))
L = [3,4,5]

B = np.zeros(A.shape)
for e in L:
    B[A==e] = 1

有更快的方法吗?

谢谢!

3 个答案:

答案 0 :(得分:3)

看起来很明显我指的是:

def AinL(A, L):
    B = np.zeros((10,), int)
    B[L] = 1
    return B[A]

基准:

 10x10 #L=3
orig     0.6665631101932377
HAL      0.4370500799268484
Psidom   1.13961720908992
PP       0.23527960386127234

 100x100 #L=3
orig     0.3015591569710523
HAL      0.29902734607458115
Psidom   0.4470538650639355
PP       0.18963343487121165

 1000x1000 #L=4
orig     0.5516874771565199
HAL      0.5967503408901393
Psidom   0.6331975681241602
PP       0.23225238709710538

 10000x1000 #L=2
orig     0.8539429588709027
HAL      0.9840140701271594
Psidom   1.0392512339167297
PP       0.7203555379528552

答案 1 :(得分:2)

以下是两个numpy选项np.in1d,它是基础python中in的矢量化版本。当数组很大时,第一个选项会显示一些加速:

选项一(快一)

np.in1d(A, L).reshape(A.shape).astype(int)

选项二(慢速)

np.apply_along_axis(np.in1d, 0, A, L).astype(int)

时序

A = np.random.randint(0, 10, (1000, 1000))
L = [3,4,5]
​
def loop():
    B = np.zeros(A.shape)
    for e in L:
        B[A==e] = 1
    return B

%timeit np.in1d(A, L).reshape(A.shape).astype(int)
# 100 loops, best of 3: 6.4 ms per loop

%timeit loop()
# 100 loops, best of 3: 16.8 ms per loop

%timeit np.apply_along_axis(np.in1d, 1, A, L).astype(int)
# 10 loops, best of 3: 21.5 ms per loop

%timeit np.apply_along_axis(np.in1d, 0, A, L).astype(int)
# 10 loops, best of 3: 35.1 ms per loop

结果检查

B1 = loop()
B2 = np.apply_along_axis(np.in1d, 0, A, L).astype(int)
B3 = np.apply_along_axis(np.in1d, 1, A, L).astype(int)
B4 = np.in1d(A, arrL).reshape(A.shape).astype(int)

(B1 == B2).all()
# True

(B1 == B3).all()
# True

(B1 == B4).all()
# True

答案 2 :(得分:2)

使用@ Psidom的1000x1000矩阵,我介绍了另外两种方法,并包含@Psidom提供的np.in1d方法。

一个使用迭代求和,另一个使用迭代按位或。

迭代按位或trial2()在下面证明,提供的结果大约比原始结果快4倍,比numpy的3d快2倍,但请注意它提供了一个矩阵结果是布尔类型。

当修改按位方法以返回整数结果trial2_int()时,其速度基本上等于numpy' s in1d

A = np.random.randint(0,10,(1000,1000))
L = [3,4,5]
def original():
  B = np.zeros(A.shape)
  for e in L:
      B[A==e] = 1
  return B

def trial1():
  B = np.empty(A.shape)
  for e in L:
    B += A == e
  return B

def trial2():
  B = A==L[0]
  for e in L[1:]:
    B |= A == e
  return B

def trial2_int():
  B = trial2()
  return B.astype(int)

def trial_Psidom():
  B = np.in1d(A,L).reshape(A.shape).astype(int)
  return B

结果:

%timeit original()
# 100 loops, best of 3: 10.5 ms per loop
%timeit trial1()
# 100 loops, best of 3: 9.43 ms per loop
%timeit trial2()
# 100 loops, best of 3: 2.37 ms per loop
%timeit trial2_int()
# 100 loops, best of 3: 5.31 ms per loop
%timeit trial_Psidom()
# 100 loops, best of 3: 5.37 ms per loop