Problem 320647 · hard · Level 03 Linear Management & Searching

Where the Second Boiler Switches On

linear regression · least squares · segmented regression · prefix sums · sorting

A greenhouse burns more gas when it is colder outside, but below some outside temperature a second boiler switches on and the relationship changes. One straight line fits the log badly; two lines, one for each side of the switching temperature, fit it well. Find the best place to split.

For a whole number s, the left part is the points with x < s and the right part the points with x >= s. A split is allowed if each part contains at least two different x values. Fit the least-squares line to each part separately; the cost of the split is the sum of the squared residuals of both lines.

Write two_lines(xs, ys) that returns a tuple (s, left, right, cost) for the allowed split with the smallest cost, where s is the smallest x of the right part, and left and right are the (intercept, slope) pairs of the two lines. If no split is allowed, return None.

The setup provides greenhouse_log(n, seed), which returns (temps, gas) for n random hours.

Examples

Input:  xs = [1, 2, 3, 4, 5, 6], ys = [2, 4, 6, 20, 21, 22]
Output: (4, (0.0, 2.0), (16.0, 1.0), 0.0)
Explanation: y = 2x fits the first three points exactly and y = 16 + x the last three.

Input:  xs = [5, 1, 3, 3, 9, 7], ys = [9, 1, 4, 6, 30, 20]
Output: (7, (-1.0, 2.0), (-15.0, 5.0), 2.0)
Explanation: the split s = 7 puts x = 1, 3, 3, 5 on the left, fitted by -1 + 2x with residuals
0, -1, 1, 0, and x = 7, 9 on the right, fitted exactly. The only other allowed split, s = 5,
costs 2.0 + 0.167. (s = 3 or s = 9 would leave one side with a single x value.)

Constraints

  • 1 <= len(xs) == len(ys) <= 10**5; all values are whole numbers between -10**4 and 10**4; the x values need not be sorted and may repeat
  • the tests have one clearly best split (every other allowed split costs noticeably more); floats are compared with a tolerance of 1e-6
  • each test must finish in well under a second in your browser: refitting both lines from scratch for every split is far too slow

Goals

  • Fit two least-squares lines, one on each side of a split point, and add up their squared errors
  • Try every split point in O(n log n) total with prefix sums over the sorted data
  • Compute sums of squared errors from sums without catastrophic cancellation
Starting Python…