Linear DPeasyCount the runs ending here, in one counter3 min · 185 of 290

Posts on an even grade

Count every stretch of fence posts whose tops step by a fixed amount, using one counter that carries all the stretches ending at a post.

A fence looks right when the post tops step by the same amount each time. The surveyor wants to know how much of the run reads that way.

The problem

A contractor drives posts along a slope and, walking the line, records the height of each post top above a fixed datum. The list heights gives those numbers in order along the fence.

A stretch is three or more consecutive posts whose heights change by the same amount from one post to the next — falling by 4 each time, rising by 12 each time, or holding level. A longer stretch contains shorter ones inside it, and each of those counts separately, because the surveyor's report lists every run of posts that reads as an even grade.

Count them all.

Input. heights — a list of integers, post heights in order along the fence.

Output. The number of stretches of three or more consecutive posts with a constant step.

Example.

heights = [4, 7, 10, 13]   ->  3

The three-post stretches [4, 7, 10] and [7, 10, 13], and the four-post one that spans the lot. All step by 3.

A second example, where a valley breaks the run:

heights = [12, 9, 6, 3, 6, 9, 12]   ->  6

Four posts fall by 3, then four posts rise by 3 — three stretches each. Nothing that straddles the bottom counts, because the step changes sign there.

Constraints.

  • 1 <= len(heights) <= 5000
  • 0 <= heights[i] <= 10^4
  • A level run counts: a step of 0 is a constant step

Hints

Hint 1

Work with the gaps between neighbouring posts rather than the heights. A stretch is a run of equal gaps.

Hint 2

Count stretches by where they end. How many stretches end at post i, given how many ended at post i - 1?

Hint 3

Chop the first post off an even stretch and what is left is still even. That is what lets one counter carry all of them.

Approach

Brute force

Take every pair of endpoints and walk between them checking the step: about n^2 / 2 stretches and up to n posts in each, so 6 x 10^10 comparisons at n = 5000. Even the tidier version that stops early is n^2 / 2 work.

The insight

Every stretch ending at post i that is longer than three is a stretch ending at post i - 1 with post i tacked on, so one running counter holds them all.

Dropping the first post of an even stretch leaves an even stretch — the step never changed. So the stretches ending at i are exactly the stretches ending at i - 1, each extended by one post, plus the brand-new three-post one. When the last two gaps agree, the counter goes up by one; when they disagree, no stretch survives and it goes back to zero.

Algorithm

  1. Set total = 0 and ending_here = 0.
  2. For each post i from index 2 onward:
  3. If heights[i] - heights[i-1] equals heights[i-1] - heights[i-2], add 1 to ending_here and add ending_here to total.
  4. Otherwise set ending_here back to 0.
  5. Return total.

Complexity

Time O(n) — one comparison per post. Space O(1) — two counters, no list of gaps needed if you compare heights directly.

Solution

Python 3 · standard library18 lines · 7 test cases, all passing
"""Posts on an even grade — count evenly stepped stretches by extending the previous one."""


def solve(heights):
    """Number of stretches of three or more consecutive posts that rise or fall by a fixed step."""
    total = 0
    ending_here = 0
    for i in range(2, len(heights)):
        if heights[i] - heights[i - 1] == heights[i - 1] - heights[i - 2]:
            # Invariant: ending_here counts the even stretches that finish at
            # post i. Extending a run of length L by one post adds exactly one
            # new stretch for every stretch that ended at the post before it,
            # plus the fresh three-post one.
            ending_here += 1
            total += ending_here
        else:
            ending_here = 0
    return total
The cases that ran
TESTS = [
    (([4, 7, 10, 13],), 3),
    (([12, 9, 6, 3, 6, 9, 12],), 6),
    (([2, 5, 9, 12, 15, 18],), 3),
    (([60, 60, 60, 60],), 3),
    (([8, 3, 9],), 0),
    (([5, 9],), 0),
    (([7],), 0),
]

Pitfalls

  • Counting only the longest stretch in each run. [4, 7, 10, 13] then gives 1 instead of 3. The report wants every stretch, nested ones included.
  • Fixing the step from the first two posts. Comparing every gap against heights[1] - heights[0] reports 3 on [12, 9, 6, 3, 6, 9, 12] instead of 6, because the second half steps the other way.
  • Leaving the counter alone when the step breaks. Without the reset, the same fence gives 10 — stretches from before the valley get extended across it.

Variants