Linear Regression (Normal Equation)

Fit an ordinary least squares linear regression using the closed-form normal equation, without calling sklearn or np.linalg.lstsq. Interviewers use this to check that you can go from the OLS objective to working matrix code and that you know the numerically sound way to solve the system.

Definition

For a feature matrix X and targets y, OLS minimizes the squared residuals ||X_aug w - y||^2, where X_aug is X with a leading column of ones so the model has an intercept. The minimizer satisfies the normal equation:

(X_aug^T X_aug) w = X_aug^T y
w = (X_aug^T X_aug)^-1 X_aug^T y

Task

Complete fit_linear_regression(X, y) so it returns the weight vector w of shape (d + 1,). Prepend a column of ones to X to model the bias, form the normal equation, and solve it with np.linalg.solve rather than an explicit matrix inverse. w[0] is the intercept and w[1:] are the coefficients for the columns of X.

Example

X = np.array([[1.0], [2.0], [3.0], [4.0]])
y = np.array([3.0, 5.0, 7.0, 9.0])   # exactly y = 2*x + 1
fit_linear_regression(X, y)          # -> array([1., 2.])

The data lies on the line y = 2x + 1, so the fit recovers bias 1.0 and slope 2.0.

Constraints

• X has shape (n, d) with n >= d + 1; y has shape (n,).
• Assume X_aug^T X_aug is invertible (full column rank).
• Use vectorized NumPy and np.linalg.solve; do not form np.linalg.inv.

Approach

Augment X with a leading column of ones so the intercept becomes just another weight, then apply the normal equation (X^T X) w = X^T y. Instead of inverting X^T X explicitly, form the two matrix products and hand the linear system to np.linalg.solve, which is faster and numerically more stable. The result is the weight vector with the bias in position 0.

Solution

import numpy as np

def fit_linear_regression(X: np.ndarray, y: np.ndarray) -> np.ndarray:
    X = np.asarray(X, dtype=float)
    y = np.asarray(y, dtype=float)

    n = X.shape[0]
    ones = np.ones((n, 1))
    X_aug = np.hstack([ones, X])          # shape (n, d + 1), bias column first

    A = X_aug.T @ X_aug                    # (d + 1, d + 1) Gram matrix
    b = X_aug.T @ y                        # (d + 1,)
    w = np.linalg.solve(A, b)             # solve A w = b, no explicit inverse
    return w

Walkthrough

On the example X = [[1], [2], [3], [4]], y = [3, 5, 7, 9]:

• X_aug = [[1, 1], [1, 2], [1, 3], [1, 4]] after prepending the ones column.
• A = X_aug^T X_aug = [[4, 10], [10, 30]] — the top-left 4 counts the samples, the off-diagonal 10 is the sum of the x values.
• b = X_aug^T y = [24, 70] — the sum of y, then the sum of x*y.
• Solving A w = b gives w = [1, 2], i.e. bias 1.0 and slope 2.0, exactly the line y = 2x + 1.

Complexity & notes

• Time O(n d^2 + d^3), space O(n d + d^2) — building X^T X costs O(n d^2), and solving the (d+1)-dimensional system costs O(d^3). For many samples and few features the O(n d^2) term dominates.
• Use np.linalg.solve(A, b) instead of np.linalg.inv(A) @ b: solving is roughly twice as fast and avoids the extra rounding error of materializing the inverse.
• The normal equation assumes X^T X is invertible (features are full column rank). If features are collinear or d > n, X^T X is singular and solve raises LinAlgError; in practice you would add L2 regularization (A + lambda*I, ridge) or fall back to np.linalg.lstsq, which uses the pseudoinverse.
• Forming X^T X squares the condition number of X, so for ill-conditioned data an SVD-based solver (lstsq) is more accurate than the normal equation.