Katkat

Implement Univariate Linear Regression with Ordinary Least Squares

Two SigmaData Scientist · Onsite

Last updated

This question evaluates understanding of univariate linear regression and the ordinary least squares (OLS) closed-form estimator, probing competencies in basic statistical estimation, interpretation of slope and intercept, numerical stability, and floating-point arithmetic.

Implement univariate (simple) linear regression fitted in batch with the ordinary least squares (OLS) closed-form solution. You must implement the formula yourself — do not use any machine-learning or statistics library fitting routine.

You are given n observations (x_1, y_1), (x_2, y_2), ..., (x_n, y_n). Find the slope b and intercept a of the line

$$\hat{y} = a + b x$$

that minimize the sum of squared residuals

$$\sum_{i=1}^{n} \left( y_i - (a + b x_i) \right)^2 .$$

The OLS closed-form solution is:

$$b = \frac{\sum_{i=1}^{n} (x_i - \bar{x})(y_i - \bar{y})}{\sum_{i=1}^{n} (x_i - \bar{x})^2}, \qquad a = \bar{y} - b,\bar{x}$$

where $\bar{x}$ and $\bar{y}$ are the means of the x and y values.

Input

  • x: a list of n real numbers (the feature values).
  • y: a list of n real numbers (the target values), where y[i] corresponds to x[i].

Output

  • A list of two numbers [b, a]: the fitted slope and intercept, in that order. Answers within an absolute or relative error of 1e-6 of the correct values are accepted.

Constraints

  • 2 <= n <= 100000
  • -10^6 <= x[i], y[i] <= 10^6
  • The x values are guaranteed to contain at least two distinct values (the denominator is never zero).
  • Your solution should run in O(n) time using a constant number of passes over the data.

Example 1

x = [1, 2, 3]
y = [2, 4, 6]

Output: [2.0, 0.0]

The points lie exactly on the line y = 2x, so the slope is 2.0 and the intercept is 0.0.

Example 2

x = [0, 1, 2, 3]
y = [1, 3, 2, 5]

Output: [1.1, 1.1]

Here $\bar{x} = 1.5$, $\bar{y} = 2.75$, $\sum (x_i - \bar{x})(y_i - \bar{y}) = 5.5$, and $\sum (x_i - \bar{x})^2 = 5.0$, giving slope b = 1.1 and intercept a = 2.75 - 1.1 * 1.5 = 1.1.

Example 3

x = [5, 5, 10]
y = [1, 3, 2]

Output: [0.0, 2.0]

The best-fit line is horizontal: the covariance term is 0, so b = 0.0 and a equals the mean of y, which is 2.0.