Nulear Rods

1921493153048_.pic_hd.jpg

import java.util.HashMap;
import java.util.Map;

public class NuclearRod {
    
    /**
     * @param n the number of rods
     * @param edges An array of Strings, pairs, like "12 23" meaning 12 and 23 is connected.
     * @return the cost 
     */
    public static int countComponents(int n, String[] edges) {
        Map roots = new HashMap<>();
        
        for (int i = 0; i < n; i++) {
            roots.put(i, i);
        }
        
        int connectedNumber = n;
        for (int i = 0; i < edges.length; i++) {
            String[] pair = edges[i].split(" ");
            int root1 = find(roots, Integer.valueOf(pair[0]));
            int root2 = find(roots, Integer.valueOf(pair[1]));
            if (root1 != root2) {
                //union root2 into root1
                roots.put(root2, root1);
                connectedNumber--;
            }
        }
        
        //calculate the number of rods in each fused part
        int[] connected = new int[n];
        for (Integer id : roots.keySet()) {
            int root = find(roots, id);
            connected[root]++;
        }
        
        //calculate the cost
        int cost = 0;
        for (int i = 0; i < connected.length; i++) {
            cost += (int)Math.ceil(Math.sqrt((double)connected[i]));
        }
        
        return cost;
    }
     
    public static int find(Map roots, int id) {
        while (roots.get(id) != id) {
            //change father to grandfather
            roots.put(id, roots.get(roots.get(id)));
            id = roots.get(id);
        }
        
        return id;
    }
    
    public static void main(String[] args) {
//        0          3
//        |          |
//        1 --- 2    4
        String[] rods = new String[] {"0 1", "1 2", "3 4"};
        int cost = countComponents(5, rods);
        System.out.println(cost); //4
    }
}