0824HardArrayHash TableTreeUnion-FindGraph TheorySorting

Level Treks Through the Huts

Tracked in this browser only

Trains the technique from

This is an original problem, written from a brief that listed the technique, the difficulty, the topics, the function shape and the input bounds — none of that problem's wording, examples, hints or editorials. The link is there so you can map your practice onto the standard set.

Same function shape, different story and different numbers.

A reserve has n shelter huts numbered 0 to n - 1. Hut i stands at height heights[i].

The huts are linked by n - 1 trails, handed over as an edge list trails, where trails[i] = [u, v] is a trail joining hut u and hut v. The trails link the whole reserve together and contain no loop, so between any two huts there is exactly one route that never visits a hut twice. That route is the trek between them.

A trek is level when both of the following hold:

the two huts at its ends stand at exactly the same height;
no hut along the route stands higher than that height.

Return how many level treks the reserve has. A single hut on its own counts as a level trek. A trek walked in one direction and the same trek walked back are one trek, counted once.

Examples

Example 1

Input: heights = [4, 9, 4], trails = [[1, 0], [1, 2]]
Output: 3

Each of the three huts on its own is a level trek. The trek between hut 0 and hut 2 has matching ends at height 4, but it runs through hut 1, which stands at height 9, so it is not level.

Example 2

Input: heights = [2, 2, 2, 2], trails = [[0, 2], [0, 1], [3, 0]]
Output: 10

Every hut stands at height 2, so no route can pass a higher hut. The four single-hut treks and the six treks between different huts are all level.

Example 3

Input: heights = [1, 1, 5, 1, 1], trails = [[0, 1], [2, 1], [2, 3], [4, 3]]
Output: 7

The five single-hut treks are level, and so are the trek from hut 0 to hut 1 and the trek from hut 3 to hut 4, both of which stay at height 1. Every other trek between two huts of height 1 runs through hut 2 at height 5.

Constraints

n == heights.length
1 <= heights.length <= 3 * 10^4
0 <= heights[i] <= 10^5
trails.length == heights.length - 1
trails[i].length == 2
0 <= trails[i][j] < heights.length, and the two huts of a trail are different
The trails form a tree: the reserve is connected and has no loop.

The signature

The editor is preloaded with this. It matches the parent problem's shape, so a solution that works here transfers to a judge unchanged.

Python

def count_level_treks(heights: list[int], trails: list[list[int]]) -> int:

Java

public int countLevelTreks(int[] heights, int[][] trails)

The intended approach

Time O(n log n) for n huts, dominated by ordering them by height.Space O(n) for the trail lists and the disjoint-set arrays.

Deal with the huts in increasing order of height, keeping the huts seen so far grouped into regions by a disjoint-set structure with iterative lookups. When a hut at height `h` is handled, it is joined to each neighbour standing at `h` or below, so a region always consists of huts no higher than the height currently being handled and the huts at that height are its tallest. Each region therefore only needs to remember its tallest height and how many of its huts reach it, and joining two regions whose tallest heights agree creates one new level trek for every pairing between those two sets of huts. Single-hut treks are added up front.

Pythonentry point: count_level_treks

def count_level_treks(heights, trails):
    # Take the huts in increasing order of height and, for each one, join it to any
    # neighbour that is no higher. Every group of huts joined together this way is a
    # region whose highest huts are exactly the ones at the height being processed,
    # so when two such regions are joined, every pairing of their highest huts is a
    # new level trek. A disjoint-set structure with iterative lookups keeps the
    # regions, and each region records its highest hut height and how many huts
    # stand at it.
    # Time: O(n log n) for n huts, dominated by the ordering by height.
    # Space: O(n) for the trail lists and the disjoint-set arrays.
    n = len(heights)
    adj = [[] for _ in range(n)]
    for pair in trails:
        adj[pair[0]].append(pair[1])
        adj[pair[1]].append(pair[0])

    owner = list(range(n))
    peak = list(heights)      # highest hut height inside the region
    at_peak = [1] * n         # how many huts of the region stand at that height

    def root_of(x):
        while owner[x] != x:
            owner[x] = owner[owner[x]]
            x = owner[x]
        return x

    order = sorted(range(n), key=lambda i: heights[i])
    treks = n                 # a single hut is a level trek on its own
    i = 0
    while i < n:
        h = heights[order[i]]
        j = i
        while j < n and heights[order[j]] == h:
            j += 1
        for k in range(i, j):
            u = order[k]
            ru = root_of(u)
            for w in adj[u]:
                if heights[w] > h:
                    continue
                rw = root_of(w)
                if rw == ru:
                    continue
                # every hut joined so far stands at h or below, so h is the peak of
                # u's region; the neighbour's region matches only if its peak is h too
                if peak[rw] == h:
                    treks += at_peak[ru] * at_peak[rw]
                    at_peak[ru] += at_peak[rw]
                owner[rw] = ru
                peak[ru] = h
        i = j
    return treks

The obvious version

Set out from each hut in turn and walk the trails with an explicit stack, never stepping onto a hut higher than the starting one. Any hut reached whose height equals the starting height, the starting hut included, closes a level trek, and only huts with an index at or above the start are counted so each trek is counted once.

This is not filler. It was written independently of the fast solution and the two were run against thousands of random inputs to confirm they agree — which is how the answer above is known to be right. Read it when the intended approach feels like a rabbit pulled out of a hat.

Python

def bf_count_level_treks(heights, trails):
    """Start from each hut in turn and walk outwards, refusing to step onto any hut
    that stands higher than the one started from. Every hut reached at exactly that
    height, itself included, ends a level trek. Walking uses an explicit stack."""
    n = len(heights)
    adj = [[] for _ in range(n)]
    for pair in trails:
        adj[pair[0]].append(pair[1])
        adj[pair[1]].append(pair[0])
    total = 0
    for start in range(n):
        ceiling = heights[start]
        seen = [False] * n
        seen[start] = True
        stack = [start]
        while stack:
            x = stack.pop()
            for y in adj[x]:
                if not seen[y] and heights[y] <= ceiling:
                    seen[y] = True
                    stack.append(y)
        for other in range(start, n):
            if seen[other] and heights[other] == ceiling:
                total += 1
    return total

How this problem was checked

Reference solution executed against 21 test cases — 21 passed
3 worked examples confirmed by running the code
0 eight-word sequences shared with the parent problem's text
Agrees with an independent brute-force implementation on random inputs
Test set proven to catch deliberately broken solutions
Reviewed by a separate quality pass
Function arity matches the parent problem's

Create a free account to unlock the solution

See the step-by-step animation, the intuition, and clean code in every language — free, no credit card.

Create a free account Already have an account? Sign in