LC 1724 - Checking Existence of Edge Length Limited Paths II

This is a copy of the solution write-up I did for this problem.

Observe that the overall length of the path is irrelevant to the solution; only the individual edge weights matter. This naturally leads us to finding the MST of the graph. This intuitively makes sense, since we greedily select edges with the lowest weight. Due to how Kruskal’s works, we will ignore the multi-edges present in the graph.

For now, I will assume that we are solving on a single tree. This assumption changes nearly nothing in the implementation.

Now that we are handling queries on a tree, the problem is much simpler. Let us build a data structure that calculates the maximum value on the path between two nodes in a tree, \( (u, v) \)

Observe that for the unique path \( u \rightarrow v \), the path can be split into two disjoint paths: the path from \( u \rightarrow LCA(u, v) \), and the path from \( v \rightarrow LCA(u, v) \). Our answer is the maximum edge along these two paths. Forcing the paths to end at \( LCA(u, v) \) allows us to easily answer queries using binary lifting.

Let \( rm \) be our range maximum “lift” array. We calculate this the exact same way that we calculate a binary lifting table, except taking the max of the two segments.

To handle queries, calculate the maximum segment while binary lifting along the path to \( LCA(u, v) \). This approach only works if both \( dep[u] = dep[v] \), so we will lift the deeper node if \( dep[u] \neq dep[v] \), lift the deeper node to the depth of the higher node, while updating the max edge value.

Observe that when handling queries, we are guaranteed a path between \( u \) and \( v \), since we are working on a single tree. Now what if the given MST is actually a minimum spanning forest? This doesn’t change anything, since our binary jumping arrays only jump to accessible nodes. The only change in the implementation is that we must DFS over every tree when initially building the data structure. Of course, it is trivial to guarantee \( u \) and \( v \) have a path between them using DSU. In fact, we can just use the same DSU that was used in Kruskal’s when calculating the MST.

The final answer is simply \( query(u, v) < limit \)

Complexity: \( O(e\log{e} + q\log{n}) \) (dropped binary lifting term as \( e + 1 \geq n \))

Code

const int maxn = 1e4 + 5;
const int lmax = 20;

struct DSU {
    vector<int> par, sz;
    DSU(int n) : par(n, -1), sz(n, 1) {}
    int find(int a) { return par[a] < 0 ? a : par[a] = find(par[a]); }
    bool join(int a, int b) {
        a = find(a), b = find(b);
        if(a == b) return false;
        if(sz[a] < sz[b]) swap(a, b);
        sz[a] += sz[b];
        par[b] = a;
        return true;
    }
};

struct LCA {
    vector<pair<int, int>> adj[maxn];
    int up[maxn][lmax + 5], rm[maxn][lmax + 5];
    int dep[maxn];
    
    void build(int n, vector<vector<int>>& edges) {
        memset(up, 0, sizeof(up));
        memset(rm, 0, sizeof(rm));
        memset(dep, 0, sizeof(dep));
        for(auto e : edges) {
            int u = e[0], v = e[1], w = e[2];
            adj[u].push_back({v, w});
            adj[v].push_back({u, w});
        }
        for(int i = 1; i <= n; i++) {
            if(up[i][0] == 0) {
                dfs(i, i, 0, 0);
            }
        }
    }
    
    void dfs(int u, int p, int w, int d) {
        up[u][0] = p;
        rm[u][0] = w;
        dep[u] = d;
        for(int l = 1; l <= 20; l++) {
            up[u][l] = up[up[u][l - 1]][l - 1];
            rm[u][l] = max(rm[u][l - 1], rm[up[u][l - 1]][l - 1]);
        }
        for(auto [v, ww] : adj[u]) {
            if(v != p) dfs(v, u, ww, d + 1);
        }
    }
    
    pair<int, int> lift(int u, int k) {
        int mx = 0;
        for(int l = lmax; l >= 0; l--) {
            if((1 << l) <= k) {
                k -= (1 << l);
                mx = max(mx, rm[u][l]);
                u = up[u][l];
            }
        }
        return {u, mx};
    }
    
    int query(int u, int v) {
        int mx = 0;
        if(dep[u] > dep[v]) swap(u, v);
        if(dep[u] < dep[v]) {
            auto [w, m] = lift(v, dep[v] - dep[u]);
            v = w;
            mx = m;
        }
        if(u == v) return mx;
        for(int l = lmax; l >= 0; l--) {
            if(up[u][l] != up[v][l]) {
                mx = max(mx, max(rm[u][l], rm[v][l]));
                u = up[u][l];
                v = up[v][l];
            }
        }
        return max(mx, max(rm[u][0], rm[v][0]));
    }
};

class DistanceLimitedPathsExist {
public:
    DSU dsu;
    LCA lca;
    
    DistanceLimitedPathsExist(int n, vector<vector<int>>& edgeList) : dsu(n + 5) {
        sort(edgeList.begin(), edgeList.end(), [](const auto &l, const auto &r) {
            return l[2] < r[2];
        });
        vector<vector<int>> mst;
        for(auto &e : edgeList) {
            if(dsu.join(++e[0], ++e[1])) {
                mst.push_back(e);
            }
        }
        lca.build(n, mst);
    }
    
    bool query(int p, int q, int limit) {
        if(dsu.find(++p) != dsu.find(++q)) return false;
        return lca.query(p, q) < limit;
    }
};

/**
 * Your DistanceLimitedPathsExist object will be instantiated and called as such:
 * DistanceLimitedPathsExist* obj = new DistanceLimitedPathsExist(n, edgeList);
 * bool param_1 = obj->query(p,q,limit);
 */