使用numpy计算,不进行迭代(矢量化)

fivyi3re  于 2023-01-05  发布在  其他
关注(0)|答案(1)|浏览(89)

我想用python/numpy矢量化(没有迭代)进行计算,得到列“D”的值。计算列“D”应该与列“A”具有相同的值,其中列“B”的值首先(在列表中向下)等于或小于列“C”的值。
在B〈= C的第一行中,设置D == A

| A | B | C | D
| 1 | 2 | 0 | 0
| 0 | 0 | 1 | 0
| 0 | 0 | 2 | 1
| 2 | 4 | 3 | 0
| 0 | 0 | 4 | 2
| 3 | 5 | 4 | 0
| 0 | 0 | 6 | 3
5gfr0r5j

5gfr0r5j1#

从文本和结果中解开逻辑,我认为这是一个总结

"""
Logic:
    if B > C:  D = 0

    The first time B <= C: D = A
    Subsequent times:
        D = previous D[ where B <= C] + 1 
"""

一种实现

import numpy as np

a = np.array( [ 1, 0, 0, 2, 0, 3, 0 ])
b = np.array( [ 2, 0, 0, 4, 0, 5, 0 ])
c = np.array( [ 0, 1, 2, 3, 4, 4, 6 ])

tgt = np.array( [ 0, 0, 1, 0, 2, 0, 3 ])

d = np.zeros_like( a )  # d == 0 unless changed

mask = b <= c
d[ mask ] = np.arange( len( d[ mask ] )) + a[ mask ][0]
# d[mask] = [ 0 1 ..              n ]    + first match in a 

d
# array([0, 0, 1, 0, 2, 0, 3])

d == tgt
# array([ True,  True,  True,  True,  True,  True,  True])

更改a数组:

a = np.array( [ 1, 7, 0, 2, 0, 3, 0 ])
d[ mask ] = np.arange( len( d[mask])) + a[ mask ][0]

d
# array([ 0,  7,  8,  0,  9,  0, 10])

相关问题