OneCompiler

42rs9xxt2

1621

Example heading with h2 size

Example heading with h3 size

Following is sample java code.
import java.util.*;

public class Main {
static final int MAXN = 100005;
static int[] sz = new int[MAXN]; // Size array for DSU
static int[] root = new int[MAXN]; // Root array for DSU

// Function to find the root of node x with path compression
static int find(int x) {
    if (root[x] == x) {
        return x;
    }
    return root[x] = find(root[x]);
}

// Function to merge two sets represented by x and y
static void merge(int x, int y) {
    x = find(x);
    y = find(y);
    if (x == y) return; // Already in the same set

    // Always merge the smaller set into the larger one
    if (sz[x] < sz[y]) {
        int temp = x;
        x = y;
        y = temp;
    }
    // Merge y into x, update the size of x
    sz[x] += sz[y];
    root[y] = x;
}

public static void main(String[] args) {
    Scanner scanner = new Scanner(System.in);
    int n = scanner.nextInt();                                                                                               
    int[][] edges = new int[n - 1][3];
    
    // Input edges: u, v, weight
    for (int i = 0; i < n - 1; i++) {
        edges[i][0] = scanner.nextInt();
        edges[i][1] = scanner.nextInt();
        edges[i][2] = scanner.nextInt();
    }
    
    // Sort edges based on weights
    Arrays.sort(edges, Comparator.comparingInt(a -> a[2]));

    // Initialize DSU
    for (int i = 1; i <= n; i++) {
        sz[i] = 1;
        root[i] = i;
    }

    long ans = 0;
    
    // Traverse sorted edges and calculate the sum
    for (int[] edge : edges) {
        int u = edge[0];
        int v = edge[1];
        int weight = edge[2];

        // Add to the answer the product of weight and the sizes of sets u and v
        ans += (long) weight * sz[find(u)] * sz[find(v)];

        // Merge the two sets
        merge(u, v);
    }
    
    // Output the final answer
    System.out.println(ans);
    scanner.close();
}

}