LC 8026. 构造乘积矩阵

这题又有点想偏了。

最开始的想法就是使用类似费马小定理的方法来解决，但是后来发现其实不行，因为费马小定理要求的是质数。虽然12345不是质数，但是它可以分解成为 `12345 = 3 x 5 x 823`. 我们可以先对所有的数进行因式分解，计算3,5,823分别有多少个，然后对于剩余下数就需要计算 `(b^-1) % MOD`. 不过这里不能使用费马小定理，因为它要求MOD是质数。但是我问了ChatGPT，如果MOD不是质数，但是b如果和MOD互质的话，那么其实可以使用扩展欧几里得算法得到 b 的逆元。

这里解法就是 `b * x = 1(% MOD)`, 我们其实可以写成 `b * x + MOD * y = 1(%MOD)`. 然后计算出x. 那么x就是 `(b^-1) % MOD`.

代码如下，整体运行时间在2s左右，题解中给出的前后缀算法只需要200ms左右。

def pow_mod(a, b, MOD):
    res = 1
    while b:
        if b & 0x1:
            res = (res * a) % MOD
        a = (a * a) % MOD
        b = b >> 1
    return res


def extended_gcd(a, b):
    if b == 0:
        return a, 1, 0
    else:
        d, x1, y1 = extended_gcd(b, a % b)
        x = y1
        y = x1 - (a // b) * y1
        return d, x, y


# 相比费马小定理，这里不要求MOD是质数，只需要确保b和MOD是互质就行
# b * x + MOD * y =  1(% MOD)
def mod_inverse(b, MOD):
    d, x, y = extended_gcd(b, MOD)
    assert (d == 1)
    return x % MOD


class Solution:
    def constructProductMatrix(self, grid: List[List[int]]) -> List[List[int]]:
        n, m = len(grid), len(grid[0])
        # 12345 = 3,5,823
        MOD = 12345
        A, B, C, D = 0, 0, 0, 1
        encode = {}
        for i in range(n):
            for j in range(m):
                a, b, c = 0, 0, 0
                x = grid[i][j]
                while x % 3 == 0:
                    a += 1
                    x = x // 3
                while x % 5 == 0:
                    b += 1
                    x = x // 5
                while x % 823 == 0:
                    c += 1
                    x = x // 823
                encode[(i, j)] = (a, b, c, x)
                A, B, C = A + a, B + b, C + c
                D = (D * x) % MOD

        # print(A, B, C, D)
        ans = [[0] * m for _ in range(n)]
        for i in range(n):
            for j in range(m):
                a, b, c, d = encode[(i, j)]
                # (D/d) % MOD
                # D % MOD * (d^(-1) % MOD)
                v = D * mod_inverse(d, MOD)
                v = (v * pow_mod(3, (A - a), MOD)) % MOD
                v = (v * pow_mod(5, (B - b), MOD)) % MOD
                v = (v * pow_mod(823, (C - c), MOD)) % MOD
                ans[i][j] = v
        return ans

题解给出的前后缀算法更加简单，只需要计算前缀乘积和后缀乘积就行，这里面就没有啥特别的数论算法了。

class Solution:
    def constructProductMatrix(self, grid: List[List[int]]) -> List[List[int]]:
        n, m = len(grid), len(grid[0])
        ans = [[0] * m for _ in range(n)]
        prev = 1
        MOD = 12345
        for i in range(n):
            for j in range(m):
                ans[i][j] = prev
                prev = prev * grid[i][j]
                prev = prev % MOD
        prev = 1
        for i in reversed(range(n)):
            for j in reversed(range(m)):
                ans[i][j] = (ans[i][j] * prev) % MOD
                prev = prev * grid[i][j]
                prev = prev % MOD
        return ans