Number of Good Paths — Hard Problem & Solution

A tree of n nodes has value vals[i] on node i. A good path is a simple path where the two endpoints carry the same value and no node strictly between them…

Difficulty: Hard
Topics: Arrays, Sorting, Graph, Union Find, Trees
Asked at: Amazon, Google, Uber
Time limit: 2 s
Memory limit: 256 MB
Languages: JavaScript, TypeScript, Python, Java, C++, C, C#, Go, Kotlin, Swift, Rust, PHP and Ruby

Problem statement

A tree of n nodes has value vals[i] on node i. A good path is a simple path where the two endpoints carry the same value and no node strictly between them carries a larger one.

Return the number of good paths. A single node is a good path of length 0, and a path and its reverse count once.

Example 1

Input: vals = [1,3,2,1,3], edges = [[0,1],[0,2],[2,3],[2,4]]
Output: 6
Explanation: The five single nodes, plus `1 → 0 → 2 → 3`.

Example 2

Input: vals = [1,1,2,2,3], edges = [[0,1],[1,2],[2,3],[2,4]]
Output: 7

Example 3

Input: vals = [1], edges = []
Output: 1
Explanation: One node is one good path.

Constraints

n == vals.length
1 <= n <= 30000
0 <= vals[i] <= 100000
edges.length == n - 1
edges[i].length == 2
0 <= edges[i][0], edges[i][1] < n
edges represents a valid tree

How to solve Number of Good Paths

Grow the tree edge by edge in increasing order of max(vals[u], vals[v]). Track, per component, its maximum value and how many nodes hold it. When two components with equal maxima join, they contribute cntA · cntB new good paths.

Approach

Start with every node its own component: best = vals[i], cnt = 1, and an answer of n for the single-node paths.
Sort the edges by the larger endpoint value and union them in that order.
If the two components' maxima are equal, add cntA · cntB and keep the combined count; otherwise keep only the larger side's maximum and count.

Why it works

The sorted order guarantees every node already merged is at most the current threshold, so nothing on the path between two maximum-valued nodes can exceed them — which is exactly the good-path condition. Pairs are counted at the moment two components merge, which is the first time such a path exists, so nothing is double counted.

Complexity

Time — O(n log n)
Space — O(n)

Pitfalls

The n single-node paths are part of the answer.
When the maxima differ, the smaller side's count is discarded, not added.
Sorting by a single endpoint rather than by the larger one lets a bigger value sneak into the middle of a path.

Reference solution

Python

from typing import List

def numberOfGoodPaths(vals: List[int], edges: List[List[int]]) -> int:
    n = len(vals)
    parent = list(range(n))

    def find(x: int) -> int:
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x

    best = vals[:]
    cnt = [1] * n
    ans = n
    for u, v in sorted(edges, key=lambda e: max(vals[e[0]], vals[e[1]])):
        ru, rv = find(u), find(v)
        if ru == rv:
            continue
        if best[ru] == best[rv]:
            ans += cnt[ru] * cnt[rv]
            new_best, new_cnt = best[ru], cnt[ru] + cnt[rv]
        elif best[ru] > best[rv]:
            new_best, new_cnt = best[ru], cnt[ru]
        else:
            new_best, new_cnt = best[rv], cnt[rv]
        parent[rv] = ru
        best[ru] = new_best
        cnt[ru] = new_cnt
    return ans

JavaScript

var numberOfGoodPaths = function(vals, edges) {
    var n = vals.length, i;
    var parent = [];
    for (i = 0; i < n; i++) parent.push(i);
    var find = function(x) {
        while (parent[x] !== x) {
            parent[x] = parent[parent[x]];
            x = parent[x];
        }
        return x;
    };
    var best = vals.slice();
    var cnt = [];
    for (i = 0; i < n; i++) cnt.push(1);
    var sorted = edges.slice();
    sorted.sort(function(a, b) {
        return Math.max(vals[a[0]], vals[a[1]]) - Math.max(vals[b[0]], vals[b[1]]);
    });
    var ans = n;
    for (i = 0; i < sorted.length; i++) {
        var ru = find(sorted[i][0]), rv = find(sorted[i][1]);
        if (ru === rv) continue;
        var newBest, newCnt;
        if (best[ru] === best[rv]) {
            ans += cnt[ru] * cnt[rv];
            newBest = best[ru];
            newCnt = cnt[ru] + cnt[rv];
        } else if (best[ru] > best[rv]) {
            newBest = best[ru];
            newCnt = cnt[ru];
        } else {
            newBest = best[rv];
            newCnt = cnt[rv];
        }
        parent[rv] = ru;
        best[ru] = newBest;
        cnt[ru] = newCnt;
    }
    return ans;
};

Also on the editorial tab: C, C#, C++, Go, Java, Kotlin, PHP, Ruby, Rust, Swift, TypeScript.

All 667 arrays problems · the whole catalogue