高效生成累积乘积三角矩阵

Question

假设我们有一个一维的向量，比如说 [a b c d]。

接下来，我们要根据这个向量构建一个特定的矩阵。

a    0   0  0
ab   b   0  0
abc  bc  c  0
abcd bcd cd d

我现在写的代码可以完成这个任务，但看起来很复杂，而且里面有一个循环，其实这个循环是完全不需要的。

import numpy as np

v = np.array([1, 2, 3])
n = len(v)
matrix = np.zeros((n,n))
for i in range(n):
    matrix [i,:i+1] = np.flip(np.cumprod(np.flip(v[:i+1])))

print(matrix)
# [[1. 0. 0.]
#  [2. 2. 0.]
#  [6. 6. 3.]]

那么，我该怎么把它变得更简单呢？

Answer 1

我不太确定你能否避免使用for循环。如果你想让代码看起来更好，可以使用 np.lib.stride_tricks.sliding_window_view 来创建你想要的步骤，然后用 cumprod 来计算乘积，最后用 np.diagflat 来构建最终下三角矩阵的一部分，每个对角线偏移量都有一部分。把这些部分加起来就能得到最终的矩阵，因为每个部分在其他元素上都是零。

import numpy as np
from numpy.lib.stride_tricks import sliding_window_view

arr = np.array([2,3,5,7])
m = sum(
    np.diagflat(sliding_window_view(arr, i+1).prod(axis=1), -i)
    for i in range(len(arr))
)

m
# returns:
array([[  2,   0,   0,   0],
       [  6,   3,   0,   0],
       [ 30,  15,   5,   0],
       [210, 105,  35,   7]])

Answer 2

抱歉，我无法处理这个请求。

Answer 3

如果你关心速度，可以考虑使用 numba：

from numba import njit

@njit
def cumprod_triangular_numba(arr):
    out = np.zeros((arr.size, arr.size), dtype=np.int64)

    for col in range(arr.size):
        p = 1
        for row in range(col, arr.size):
            p *= arr[row]
            out[row, col] = p

    return out

---

基准测试：

import numpy as np
import perfplot
from numba import njit, prange
from numpy.lib.stride_tricks import sliding_window_view


def cumprod_triangular_orig(arr):
    n = len(arr)
    matrix = np.zeros((n, n))
    for i in range(n):
        matrix[i, : i + 1] = np.flip(np.cumprod(np.flip(arr[: i + 1])))
    return matrix


def cumprod_triangular_james(arr):
    return sum(
        np.diagflat(sliding_window_view(arr, i + 1).prod(axis=1), -i)
        for i in range(len(arr))
    )


def cumprod_triangular_onyambu_1(arr):
    u = arr.cumprod()
    return u[:, None] / np.r_[1, u[:-1]] * np.tri(arr.size, dtype=int)


def cumprod_triangular_onyambu_2(arr):
    a = np.triu(arr).T
    i1 = a == 0
    a[i1] = 1
    return np.where(i1, 0, a.cumprod(0))


def cumprod_triangular_onyambu_3(arr):
    a = np.triu(arr).T
    return np.where(a, a, 1).cumprod(0) * np.tri(arr.size, dtype=int)


@njit
def cumprod_triangular_numba(arr):
    out = np.zeros((arr.size, arr.size), dtype=np.int64)

    for col in prange(arr.size):
        p = 1
        for row in range(col, arr.size):
            p *= arr[row]
            out[row, col] = p

    return out


@njit(parallel=True)
def cumprod_triangular_numba_parallel(arr):
    out = np.zeros((arr.size, arr.size), dtype=np.int64)

    for col in prange(arr.size):
        p = 1
        for row in range(col, arr.size):
            p *= arr[row]
            out[row, col] = p

    return out


arr = np.array([2, 3, 5, 7])

assert np.allclose(cumprod_triangular_numba(arr), cumprod_triangular_orig(arr))
assert np.allclose(cumprod_triangular_numba_parallel(arr), cumprod_triangular_orig(arr))
assert np.allclose(cumprod_triangular_james(arr), cumprod_triangular_orig(arr))
assert np.allclose(cumprod_triangular_onyambu_1(arr), cumprod_triangular_orig(arr))
assert np.allclose(cumprod_triangular_onyambu_2(arr), cumprod_triangular_orig(arr))
assert np.allclose(cumprod_triangular_onyambu_3(arr), cumprod_triangular_orig(arr))


np.random.seed(0)

perfplot.show(
    setup=lambda n: np.random.randint(1, 2, size=n, dtype=np.int64),
    kernels=[
        cumprod_triangular_orig,
        cumprod_triangular_james,
        cumprod_triangular_numba,
        cumprod_triangular_numba_parallel,
        cumprod_triangular_onyambu_1,
        cumprod_triangular_onyambu_2,
        cumprod_triangular_onyambu_3,
    ],
    labels=["orig", "james", "numba", "numba_parallel", "o_1", "o_2", "o_3"],
    n_range=[3, 5, 10, 15, 20, 25, 50, 75, 100, 200, 500, 1000],
    xlabel="N",
    logx=True,
    logy=True,
    equality_check=np.allclose,
)

在我的电脑上（AMD 5700x），生成了这个图：