Asked at Two Sigma, Jump Trading, Squarepoint Capital
Theory: Quant Python That Survives ReviewRead the problem, hints and solution here. The editor needs a bigger screen: open this page on a laptop to write and run your code.
"Implement linear regression in numpy" is one of the most consistently reported coding tasks in quant researcher online assessments. Not because firms need another regression library: they want to see whether you know what the fit actually computes when the convenience wrapper is taken away.
Implement ols_fit(X, y).
X is a list of feature rows (each a list of floats), without an
intercept column. It may also be a flat list of floats for the
single-feature case.y is a list of floats, one target per row.Fit ordinary least squares with an intercept, and return a tuple
(coeffs, r2):
coeffs: the fitted coefficients as a list of floats rounded to 4
decimals, intercept first, then one coefficient per feature in input
order.r2: the coefficient of determination \( R^2 = 1 - SS_{res}/SS_{tot} \),
rounded to 4 decimals. If y is constant, define \( R^2 = 1.0 \).Columns may be collinear. Solve by least squares (numpy.linalg.lstsq),
which returns the minimum-norm solution in the degenerate case; the normal
equations with a plain inverse will crash there.
ols_fit([[1], [2], [3]], [3, 5, 7])
# ([1.0, 2.0], 1.0)
ols_fit([[1, 2], [2, 1], [3, 4], [4, 3]], [10, 11, 20, 21])
# ([3.0, 3.0, 2.0], 1.0)
A hidden test runs 20,000 rows by 8 features, which any vectorised solution handles instantly.