给定一个由大小为nxm的整数a和一个整数b组成的矩阵。在给定的矩阵中，每一行和每一列都按递增顺序排序。以给定形式查找并返回b在矩阵中的位置：
如果a[i][j]=b，则返回（i*1009+j）
如果b不存在，返回-1。
我的解决方案：

public static int solve(int[][] A, int B) {
    int N = A.length;
    int i = 0;
    int j = N - 1;
    while (i < N && j >= 0) {
        if (A[i][j] == B) {
            System.out.println("number found");
            return (i * 1009 + j);
        }
        if (A[i][j] > B) {
            j--;
        } else {
            i++;
        }
    }
    System.out.println("number not found");
    return (-1);
}

当矩阵不是nxn时不起作用。我怎样才能修好它？