如何矢量化/优化依赖于先前行的计算

Question

I'm working on something where runtime is extremely important and the data we're working with is large, but basically the problem boils down to optimizing solving for series x where x1 is known and x = ax+b from the prior row. So for example, starting state:

a b x
1 2 3
3 1
2 2
4 8
1 9

and the end state would look like this:

a b x
1 2 3
3 1 5
2 2 16
4 8 72
1 9 81

because 3*1+2 = 5, 5*3+1 = 16, etc.

I tried working out the math of it and I ended up with:

b0 = x1
xi = sum(n=0 to i-1)(bn*product(m=n+1 to i-1)(am)

So for example for the 3rd row you'd end up with:

x3 = a1*a2*b0 + b1*a2 + b2 = 3*1*3 + 2*3 + 1 = 9 + 6 + 1 = 16

But computationally that seems to be worse than just calculating each x by looping over rows, something like this:

for i in range(2,len(df)):
    df.x[i] = df.x[i-1]*df.a[i-1]+df.b[i-1]

Is there an easier way to tackle this that I'm missing, or am I just dealing with a computationally expensive operation that I'll have to eat the cost of iterating? If the a term didn't exist the bn portion could be tackled via cumsum, something like:

df['b_cumsum'] = x1+cumsum(df.b)

but I end up hitting a wall when trying to include the a terms, especially since we end up needing so many different sets of products even within each sum term.

Thanks.

Answer 1

When I run into functions I cannot vectorize, but it needs to be efficient, I use numba . Which is a just in time compilation (JIT) module. Most of the times this can be even faster than the native pandas methods:

from numba import njit

@njit
def calculation(arr):
    result = np.empty(arr.shape[0])
    for idx, row in enumerate(arr):
        if idx == 0:
            result[idx] = row[2]
        else:
            row = arr[idx-1]
            result[idx] = result[idx-1] * row[0] + row[1]
    
    return result

df['x'] = calculation(df.to_numpy())

print(df)

   a  b      x
0  1  2    3.0
1  3  1    5.0
2  2  2   16.0
3  4  8   34.0
4  1  9  144.0

---

note : when you want to time it. Don't time it on the first run, since it did not compile yet. First run it once, then time it on the second run.

Answer 2

You can first calculate a rescaled x : x ' = x /cumprod( a ) using a matching b ' = b /cumprod( a )

This can be done with vectorized operations as can the backtransform from x ' to x :

ab = np.array([[1, 2],
               [3, 1],
               [2, 2],
               [4, 8],
               [1, 9]])

scale = ab.T[0].cumprod()
xp = 3+(ab.T[1]/scale).cumsum()
x = xp*scale
x
array([  5.,  16.,  34., 144., 153.])