Implement Univariate Linear Regression with Ordinary Least Squares
Two SigmaData Scientist · OnsiteLast 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 ofnreal numbers (the feature values).y: a list ofnreal numbers (the target values), wherey[i]corresponds tox[i].
Output
- A list of two numbers
[b, a]: the fitted slope and intercept, in that order. Answers within an absolute or relative error of1e-6of the correct values are accepted.
Constraints
2 <= n <= 100000-10^6 <= x[i], y[i] <= 10^6- The
xvalues 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.
Solution walkthrough is for members
Get the full walkthrough with reference code and complexity analysis, plus Kitty Coach at 300 messages a day. Hints, the code runner and community answers stay free.