Minimum operations to make all the edges of a path equal

Given a tree with N nodes and (N – 1) edges in the form of edges[], such that edges[i] = {u, v, w} denotes an edge from node u to node v with weight w. We are also given Q queries such that queries[i] = {node1, node2} and the task is to find the minimum number of operations to make the weight of all edges between node1 and node2 equal. In one operation, we can change the weight of any edge to any other value. The root of the tree is node 0, so par[0] = -1. The queries are independent of each other.

Constraints:
1 <= N <= 10^4
1 <= wi <= 26

Examples:

Input: N = 4, edges[][] = {{0, 1, 1}, {0, 2, 1}, {1, 3, 2}}, queries[][] = {{2, 3}, {0, 2}}
Output: 1 0
Explanation:

• Query 1: There are 2 edges of weight 1 and 1 edge of weight 2 between node 2 and 3. So, we can change the weight of the edge whose weight is 2 and update it to 1. Now, all the edges have same weight, that is 1.
• Query 2: There is only 1 edge with weight 1 between node 0 and node 2. So, no need to apply any operations.

Input: N = 5, edges[][] = {{0, 1, 2}, {1, 2, 2}, {2, 3, 2}, {3, 4, 1}}, queries[][] = {{0, 4}}
Output: 1
Explanation:

• Query 1: There are 3 edges of weight 2 and 1 edge of weight 1 between node 0 and 4. So, we can change the weight of the edge whose weight is 1 and update it to 2. Now, all the edges have same weight, that is 2.

Approach: To solve the problem, follow the below idea:

The problem can be solved using Binary Lifting. For each query(u, v), to minimize the number of operations to make all the edges equal we need to find the most frequent weight among all the edges and convert the weights of all edges equal to the most frequent weight. This can be done by finding the LCA of node u and v. Since the range of weights given is very small (<=26), we can check for all the weights from 1 to 26.

Let’s say we want to make all the edges from node u to v to be equal to weight X. So, we have two problems:

• How to count the number of edges from node u to v?
• How to count the number of edges already having the weight of X?

For the first problem, we can find the depth of all nodes and then the number of nodes between node u and v can be calculated as: depth[u] – depth[LCA(u, v)] + depth[v] – depth[LCA(u, v)] = depth[u] + depth[v] – depth[LCA(u, v)]

For the second problem, we can maintain a 2D array to store prefix sums, say pref such that pref[i][j] stores the number of edges with weight i from root to node j, that is pref[i][j] stores the prefix count of weight i till node j. Now, we can calculate the number of edges between node u and v that already have a weight of X as pref[X][u] – pref[X][LCA(u,v)] + pref[X][v] – pref[X][LCA(u, v)] = pref[X][u] + pref[X][v] – 2 * pref[X][LCA(u, v)].

Step-by-step algorithm:

• Initialize a depth array, say depth[] such that depth[i] stores the depth of ith node.
• Initialize a 2D array parent[][] such that parent[i][j] stores the (2^j)th parent of node i.
• Initialize a prefix count array, say pref[][] such that pref[i][j] stores the number of edges with weight i from root to node j.
• Run a DFS to initialize the depth[] array, the pref[][] array and the first column of the parent[][] array.
• Fill the remaining columns of the parent[][] array using previous column, parent[i][j] = parent[parent[i][j-1]][j-1].
• For each query, find the count of the most frequent edge weight using the parent[][] and pref[][] array and subtract it from the total number of edges between node u and node v to get the answer.

Below is the implementation of algorithm:

C++

#include <bits/stdc++.h>
using namespace std;

// DFS function to initialize the depth[] array, the
// pref[][] array and the first column of the parent[][]
// array.
void dfs(int node, int parentNode, int d,
        vector<int>& depth,
        vector<vector<pair<int, int> > >& graph,
        vector<vector<int> >& parent,
        vector<vector<int> >& pref)
{
    parent[0][node] = parentNode;
    depth[node] = d;
    for (auto ele : graph[node]) {
        int child = ele.first;
        int weight = ele.second;
        if (child == parentNode)
            continue;
        pref[child] = pref[node];
        pref[child][weight]++;
        dfs(child, node, d + 1, depth, graph, parent, pref);
    }
}

// Function to find the LCA of node u and v
int lca(int u, int v, vector<int>& depth,
        vector<vector<int> >& parent, int m)
{
    if (depth[u] > depth[v])
        swap(u, v);

    // adjust x and y to the same depth
    for (int p = 0; (1 << p) <= depth[v] - depth[u]; ++p)
        if ((1 << p) & depth[v] - depth[u])
            v = parent[p][v];

    // jump u and v together if they have different
    // ancestors
    for (int p = m - 1; p >= 0; --p) {
        if (parent[p][u] != parent[p][v]) {
            u = parent[p][u];
            v = parent[p][v];
        }
    }
    return u == v ? u : parent[0][u];
}

vector<int> solve(int N, vector<vector<int> >& edges,
                vector<vector<int> >& queries)
{
    int m = (int)log2(N) + 1, C = 27;
    vector<vector<pair<int, int> > > graph(N);
    for (auto& edge : edges) {
        graph[edge[0]].push_back({ edge[1], edge[2] });
        graph[edge[1]].push_back({ edge[0], edge[2] });
    }

    // parent[i][j] means the 2^i -th father of vertex j:
    vector<vector<int> > parent(m, vector<int>(N));
    // pref[i][j] means the count of j from root to vertex
    // i:
    vector<vector<int> > pref(N);
    // depth[i] means the depth of vertex i:
    vector<int> depth(N, 0);

    pref[0] = vector<int>(C, 0);
    dfs(0, 0, 0, depth, graph, parent, pref);

    // binary lifting
    for (int i = 1; i < m; ++i) {
        for (int j = 0; j < N; ++j) {
            parent[i][j] = parent[i - 1][parent[i - 1][j]];
        }
    }

    vector<int> res;
    for (auto& q : queries) {
        int x = q[0], y = q[1],
            l = lca(x, y, depth, parent, m);

        // the total length between x and y:
        int len = depth[x] + depth[y] - 2 * depth[l];

        // the most frequent weight between x and y:
        int X = 0;

        for (int z = 1; z < C; ++z) {
            int currentX
                = pref[x][z] + pref[y][z] - pref[l][z] * 2;
            X = max(X, currentX);
        }

        // the others must be changed:
        res.push_back(len - X);
    }
    return res;
}

int main()
{
    int N = 5;
    vector<vector<int> > queries = { { 0, 4 } };
    vector<vector<int> > edges = {
        { 0, 1, 2 }, { 1, 2, 2 }, { 2, 3, 2 }, { 3, 4, 1 }
    };
    vector<int> ans = solve(N, edges, queries);

    for (int i = 0; i < ans.size(); i++) {
        cout << ans[i] << " ";
    }
    cout << endl;
    return 0;
}

Java

import java.util.*;

public class Main {

    // DFS function to initialize the depth[] array, the
    // pref[][] array and the first column of the parent[][]
    // array.
    public static void dfs(int node, int parentNode, int d,
                           List<Integer> depth,
                           List<List<Map.Entry<Integer, Integer>>> graph,
                           List<List<Integer>> parent,
                           List<List<Integer>> pref) {
        parent.get(0).set(node, parentNode);
        depth.set(node, d);
        for (Map.Entry<Integer, Integer> ele : graph.get(node)) {
            int child = ele.getKey();
            int weight = ele.getValue();
            if (child == parentNode)
                continue;
            pref.set(child, new ArrayList<>(pref.get(node)));
            pref.get(child).set(weight, pref.get(child).get(weight) + 1);
            dfs(child, node, d + 1, depth, graph, parent, pref);
        }
    }

    // Function to find the LCA of node u and v
    public static int lca(int u, int v, List<Integer> depth,
                          List<List<Integer>> parent, int m) {
        if (depth.get(u) > depth.get(v))
            u = v + (v = u) - u; // swapping using xor trick

        // adjust x and y to the same depth
        for (int p = 0; (1 << p) <= depth.get(v) - depth.get(u); ++p)
            if (((1 << p) & (depth.get(v) - depth.get(u))) != 0)
                v = parent.get(p).get(v);

        // jump u and v together if they have different
        // ancestors
        for (int p = m - 1; p >= 0; --p) {
            if (!Objects.equals(parent.get(p).get(u), parent.get(p).get(v))) {
                u = parent.get(p).get(u);
                v = parent.get(p).get(v);
            }
        }
        return Objects.equals(u, v) ? u : parent.get(0).get(u);
    }

    public static List<Integer> solve(int N, List<List<Integer>> edges,
                                       List<List<Integer>> queries) {
        int m = (int) (Math.log(N) / Math.log(2)) + 1, C = 27;
        List<List<Map.Entry<Integer, Integer>>> graph = new ArrayList<>();
        for (int i = 0; i < N; i++)
            graph.add(new ArrayList<>());
        for (List<Integer> edge : edges) {
            int u = edge.get(0), v = edge.get(1), w = edge.get(2);
            graph.get(u).add(new AbstractMap.SimpleEntry<>(v, w));
            graph.get(v).add(new AbstractMap.SimpleEntry<>(u, w));
        }

        // parent[i][j] means the 2^i -th father of vertex j:
        List<List<Integer>> parent = new ArrayList<>();
        for (int i = 0; i < m; i++)
            parent.add(new ArrayList<>(Collections.nCopies(N, 0)));
        // pref[i][j] means the count of j from root to vertex
        // i:
        List<List<Integer>> pref = new ArrayList<>();
        for (int i = 0; i < N; i++)
            pref.add(new ArrayList<>(Collections.nCopies(C, 0)));
        // depth[i] means the depth of vertex i:
        List<Integer> depth = new ArrayList<>(Collections.nCopies(N, 0));

        pref.set(0, new ArrayList<>(Collections.nCopies(C, 0)));
        dfs(0, 0, 0, depth, graph, parent, pref);

        // binary lifting
        for (int i = 1; i < m; ++i) {
            for (int j = 0; j < N; ++j) {
                parent.get(i).set(j, parent.get(i - 1).get(parent.get(i - 1).get(j)));
            }
        }

        List<Integer> res = new ArrayList<>();
        for (List<Integer> q : queries) {
            int x = q.get(0), y = q.get(1),
                    l = lca(x, y, depth, parent, m);

            // the total length between x and y:
            int len = depth.get(x) + depth.get(y) - 2 * depth.get(l);

            // the most frequent weight between x and y:
            int X = 0;

            for (int z = 1; z < C; ++z) {
                int currentX
                        = pref.get(x).get(z) + pref.get(y).get(z) - pref.get(l).get(z) * 2;
                X = Math.max(X, currentX);
            }

            // the others must be changed:
            res.add(len - X);
        }
        return res;
    }

    public static void main(String[] args) {
        int N = 5;
        List<List<Integer>> queries = new ArrayList<>();
        queries.add(Arrays.asList(0, 4));
        List<List<Integer>> edges = Arrays.asList(
                Arrays.asList(0, 1, 2),
                Arrays.asList(1, 2, 2),
                Arrays.asList(2, 3, 2),
                Arrays.asList(3, 4, 1)
        );
        List<Integer> ans = solve(N, edges, queries);

        for (int a : ans) {
            System.out.print(a + " ");
        }
        System.out.println();
    }
}

C#

using System;
using System.Collections.Generic;
using System.Linq;

public class GFG
{
    // DFS function to initialize the depth[] array, the
    // pref[][] array, and the first column of the parent[][]
    // array.
    public static void dfs(int node, int parentNode, int d,
                        List<int> depth,
                        List<List<KeyValuePair<int, int>>> graph,
                        List<List<int>> parent,
                        List<List<int>> pref)
    {
        parent[0][node] = parentNode;
        depth[node] = d;
        foreach (var ele in graph[node])
        {
            int child = ele.Key;
            int weight = ele.Value;
            if (child == parentNode)
                continue;
            pref[child] = new List<int>(pref[node]);
            pref[child][weight] += 1;
            dfs(child, node, d + 1, depth, graph, parent, pref);
        }
    }

    // Function to find the LCA of node u and v
    public static int lca(int u, int v, List<int> depth,
                        List<List<int>> parent, int m)
    {
        if (depth[u] > depth[v])
            u = v + (v = u) - u; // swapping using xor trick

        // adjust x and y to the same depth
        for (int p = 0; (1 << p) <= depth[v] - depth[u]; ++p)
            if (((1 << p) & (depth[v] - depth[u])) != 0)
                v = parent[p][v];

        // jump u and v together if they have different
        // ancestors
        for (int p = m - 1; p >= 0; --p)
        {
            if (parent[p][u] != parent[p][v])
            {
                u = parent[p][u];
                v = parent[p][v];
            }
        }
        return (u == v) ? u : parent[0][u];
    }

    public static List<int> solve(int N, List<List<int>> edges,
                                List<List<int>> queries)
    {
        int m = (int)(Math.Log(N) / Math.Log(2)) + 1;
        int C = 27;

        List<List<KeyValuePair<int, int>>> graph = new List<List<KeyValuePair<int, int>>>();
        for (int i = 0; i < N; i++)
            graph.Add(new List<KeyValuePair<int, int>>());
        foreach (List<int> edge in edges)
        {
            int u = edge[0], v = edge[1], w = edge[2];
            graph[u].Add(new KeyValuePair<int, int>(v, w));
            graph[v].Add(new KeyValuePair<int, int>(u, w));
        }

        // parent[i][j] means the 2^i -th father of vertex j:
        List<List<int>> parent = new List<List<int>>();
        for (int i = 0; i < m; i++)
            parent.Add(Enumerable.Repeat(0, N).ToList());
        // pref[i][j] means the count of j from root to vertex
        // i:
        List<List<int>> pref = new List<List<int>>();
        for (int i = 0; i < N; i++)
            pref.Add(Enumerable.Repeat(0, C).ToList());
        // depth[i] means the depth of vertex i:
        List<int> depth = Enumerable.Repeat(0, N).ToList();

        pref[0] = Enumerable.Repeat(0, C).ToList();
        dfs(0, 0, 0, depth, graph, parent, pref);

        // binary lifting
        for (int i = 1; i < m; ++i)
        {
            for (int j = 0; j < N; ++j)
            {
                parent[i][j] = parent[i - 1][parent[i - 1][j]];
            }
        }

        List<int> res = new List<int>();
        foreach (List<int> q in queries)
        {
            int x = q[0], y = q[1],
                    l = lca(x, y, depth, parent, m);

            // the total length between x and y:
            int len = depth[x] + depth[y] - 2 * depth[l];

            // the most frequent weight between x and y:
            int X = 0;

            for (int z = 1; z < C; ++z)
            {
                int currentX
                        = pref[x][z] + pref[y][z] - pref[l][z] * 2;
                X = Math.Max(X, currentX);
            }

            // the others must be changed:
            res.Add(len - X);
        }
        return res;
    }

    public static void Main(string[] args)
    {
        int N = 5;
        List<List<int>> queries = new List<List<int>>();
        queries.Add(new List<int> { 0, 4 });
        List<List<int>> edges = new List<List<int>>
        {
            new List<int> { 0, 1, 2 },
            new List<int> { 1, 2, 2 },
            new List<int> { 2, 3, 2 },
            new List<int> { 3, 4, 1 }
        };
        List<int> ans = solve(N, edges, queries);

        foreach (int a in ans)
        {
            Console.Write(a + " ");
        }
        Console.WriteLine();
    }
}

JavaScript

// DFS function to initialize the depth[] array, the
// pref[][] array and the first column of the parent[][]
// array.
function dfs(node, parentNode, d, depth, graph, parent, pref) {
    parent[0][node] = parentNode;
    depth[node] = d;
    for (let ele of graph[node]) {
        let child = ele[0];
        let weight = ele[1];
        if (child === parentNode)
            continue;
        pref[child] = pref[node].slice();
        pref[child][weight]++;
        dfs(child, node, d + 1, depth, graph, parent, pref);
    }
}

// Function to find the LCA of node u and v
function lca(u, v, depth, parent, m) {
    if (depth[u] > depth[v])
        [u, v] = [v, u];

    // adjust x and y to the same depth
    for (let p = 0; (1 << p) <= depth[v] - depth[u]; ++p)
        if ((1 << p) & (depth[v] - depth[u]))
            v = parent[p][v];

    // jump u and v together if they have different
    // ancestors
    for (let p = m - 1; p >= 0; --p) {
        if (parent[p][u] !== parent[p][v]) {
            u = parent[p][u];
            v = parent[p][v];
        }
    }
    return u === v ? u : parent[0][u];
}

function solve(N, edges, queries) {
    let m = Math.floor(Math.log2(N)) + 1, C = 27;
    let graph = Array.from({length: N}, () => []);
    for (let edge of edges) {
        graph[edge[0]].push([edge[1], edge[2]]);
        graph[edge[1]].push([edge[0], edge[2]]);
    }

    // parent[i][j] means the 2^i -th father of vertex j:
    let parent = Array.from({length: m}, () => Array(N).fill(0));
    // pref[i][j] means the count of j from root to vertex
    // i:
    let pref = Array.from({length: N}, () => Array(C).fill(0));
    // depth[i] means the depth of vertex i:
    let depth = Array(N).fill(0);

    dfs(0, 0, 0, depth, graph, parent, pref);

    // binary lifting
    for (let i = 1; i < m; ++i) {
        for (let j = 0; j < N; ++j) {
            parent[i][j] = parent[i - 1][parent[i - 1][j]];
        }
    }

    let res = [];
    for (let q of queries) {
        let x = q[0], y = q[1],
            l = lca(x, y, depth, parent, m);

        // the total length between x and y:
        let len = depth[x] + depth[y] - 2 * depth[l];

        // the most frequent weight between x and y:
        let X = 0;

        for (let z = 1; z < C; ++z) {
            let currentX = pref[x][z] + pref[y][z] - pref[l][z] * 2;
            X = Math.max(X, currentX);
        }

        // the others must be changed:
        res.push(len - X);
    }
    return res;
}

let N = 5;
let queries = [[0, 4]];
let edges = [[0, 1, 2], [1, 2, 2], [2, 3, 2], [3, 4, 1]];
let ans = solve(N, edges, queries);

for (let i = 0; i < ans.length; i++) {
    console.log(ans[i] + " ");
}
console.log("\n");

Python3

from collections import defaultdict
from math import log2
# DFS function to initialize the depth[], pref[][] arrays,
# and the first column of the parent[][] array.
def dfs(node, parentNode, d, depth, graph, parent, pref):
    parent[0][node] = parentNode
    depth[node] = d
    for child, weight in graph[node]:
        if child == parentNode:
            continue
        pref[child] = pref[node][:]
        pref[child][weight] += 1
        dfs(child, node, d + 1, depth, graph, parent, pref)

# Function to find the LCA of node u and v
def lca(u, v, depth, parent, m):
    if depth[u] > depth[v]:
        u, v = v, u

    # Adjust u and v to the same depth
    for p in range(m):
        if (1 << p) & (depth[v] - depth[u]):
            v = parent[p][v]

    # Jump u and v together if they have different ancestors
    for p in range(m - 1, -1, -1):
        if parent[p][u] != parent[p][v]:
            u = parent[p][u]
            v = parent[p][v]

    return u if u == v else parent[0][u]

def solve(N, edges, queries):
    m = int(log2(N)) + 1
    C = 27

    graph = defaultdict(list)
    for edge in edges:
        graph[edge[0]].append((edge[1], edge[2]))
        graph[edge[1]].append((edge[0], edge[2]))

    parent = [[0] * N for _ in range(m)]
    pref = [[] for _ in range(N)]
    depth = [0] * N

    pref[0] = [0] * C
    dfs(0, 0, 0, depth, graph, parent, pref)

    # Binary lifting
    for i in range(1, m):
        for j in range(N):
            parent[i][j] = parent[i - 1][parent[i - 1][j]]

    res = []
    for q in queries:
        x, y = q[0], q[1]
        l = lca(x, y, depth, parent, m)

        # Total length between x and y
        len_ = depth[x] + depth[y] - 2 * depth[l]
        X = 0

        # The most frequent weight between x and y
        for z in range(1, C):
            currentX = pref[x][z] + pref[y][z] - pref[l][z] * 2
            X = max(X, currentX)

        # The others must be changed
        res.append(len_ - X)

    return res

if __name__ == "__main__":
    N = 5
    queries = [[0, 4]]
    edges = [[0, 1, 2], [1, 2, 2], [2, 3, 2], [3, 4, 1]]
    ans = solve(N, edges, queries)

    for i in range(len(ans)):
        print(ans[i], end=" ")
    print()
    
# This code is contributed by shivamgupta310570

Output

1 

Time Complexity: O(N * logN + Q * logN), where N is the number of nodes and Q is the number of queries.
Auxiliary Space: O(N * logN)