给我一段求解Helmoltz方程direchlet问题的python代码

以下是一段求解Helmholtz方程direchlet问题的Python代码：

import numpy as np
from scipy.sparse import diags
from scipy.sparse.linalg import spsolve

def helmholtz_dirichlet(N, k, f, g):
    h = 1 / (N + 1)
    x = np.linspace(h, 1 - h, N)
    y = np.linspace(h, 1 - h, N)
    X, Y = np.meshgrid(x, y)
    A = diags([-4*np.ones(N**2), np.ones(N**2-1), np.ones(N**2-1), np.ones(N**2-N), np.ones(N**2-N)], [0, 1, -1, N, -N])
    A = A / h**2 + k**2 * np.eye(N**2)
    b = f(X, Y)
    b[0, :] += g(x)
    b[-1, :] += g(x)
    b[:, 0] += g(y)
    b[:, -1] += g(y)
    b = b.flatten()
    u = spsolve(A, b)
    return u.reshape((N, N))

其中，N是网格的大小，k是波数，f是右端项函数，g是边界条件函数。函数返回一个N×N的矩阵，表示解的近似值。