使用numba模块优化python代码

Question

我有这样的问题：

” 在游戏中，现代地图由 n x n 的正方形网格表示。表的行从上到下从 1 到 n 编号。表格的列从左到右从 1 到 n 编号。位于表的第 i 行和第 j 列交叉点的单元格称为单元格 (i, j)。 单元格 (i,j) 包含一个整数 r，表示半径为 r 的风暴，风暴中心位于单元格 (i, j)；当 |x − i| 时，地块 (x,y) 将受到地块 (i,j) 处风暴的影响+ |y - j| ≤ r（曼哈顿距离）。如果 r = 0，则该单元格中没有风暴。

需求：统计没有受到风暴影响的小区数量

Data：来自标准输入设备的输入包括： 第一行包含正整数 n (n < 3000) describing the size of the table; The next n lines each contain n integers r (0

结果：向标准输出写入一个数字，即未受风暴影响的单元数。 “

我写了一段代码来解决这个问题。虽然输出是正确的，但当尝试输入大数字时，运行时间很长（这与我的硬件配置无关）。

我的代码：

import numpy as np
from scipy.optimize import minimize

def find_three_numbers(A, N, T):
    A = np.array(A)

    def objective(x):
        avg = np.mean(x)
        diff = abs(avg - T)

        return -diff

    def constraint(x):
        sq_diff = np.sum((x - A) ** 2)

        return -sq_diff

    cons = {'type': 'ineq', 'fun': constraint}

    x0 = A[:3]

    result = minimize(objective, x0, constraints=cons)

    if result.success:
        return np.sum(result.x)
    else:
        raise ValueError(result.message)

def main():
    N, T = map(int, input("Input integers N and T: ").split())

    A = list(map(int, input("Input array A: ").split()))

    result = find_three_numbers(A, N, T)
    print("Result:", result)

if __name__ == "__main__":
    main()

我看了很多不同的权威文章，但我仍然无法优化我的代码。我最适合使用 numba。

希望大家帮忙。非常感谢！

P/s：如有错误，希望大家谅解。再次非常感谢！ :D