Weighted Least Squares Diagonal

Problem

Solve weighted least squares min ||W^(1/2)(A*x - b)||^2 for A of shape n x 2 with diagonal weight matrix W = diag(w).

Build the normal equations (A^T W A) x = A^T W b and solve the 2x2 system. Return [x0, x1], or None if any weight is negative or A^T W A is singular.

Hint

Weights enter both sides of the normal equations; zero weights ignore those rows entirely.

case 1

in: [[[1,0],[0,1],[1,1]],[1,2,3],[1,1,1]]

expected: [1,2]

case 2

in: [[[1,1],[1,2],[1,3]],[1,2,3],[1,1,1]]

expected: [0,1]

case 3

in: [[[1,1],[2,1],[3,1]],[2,3,4],[2,1,2]]

expected: [1,1]

case 4

in: [[[1,0],[1,0]],[1,2],[1,1]]

expected: null

case 5

in: [[[1,0],[0,1]],[1,1],[-1,1]]

expected: null

Starter codePython

def weighted_least_squares(A, b, w):
    # Your code here
    pass

5 test cases · run your code against them in the editor