[LintCode] Problem 629 - Minimum Spanning Tree

Posted on | In ,

Given a list of Connections, which is the Connection class (the city name at both ends of the edge and a cost between them), find some edges, connect all the cities and spend the least amount.

Return the connects if can connect all the cities, otherwise return empty list.

Note

Return the connections sorted by the cost, or sorted city1 name if their cost is same, or sorted city2 if their city1 name is also same.

Example

Gievn the connections = [“Acity”,”Bcity”,1], [“Acity”,”Ccity”,2], [“Bcity”,”Ccity”,3]

Return [“Acity”,”Bcity”,1], [“Acity”,”Ccity”,2]

Code

1
2
3
4
5
6
7
8
9
public class Connection {
    public String city1, city2;
    public int cost;
    public Connection(String city1, String city2, int cost) {
        this.city1 = city1;
        this.city2 = city2;
        this.cost = cost;
    }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
public class UnionFind{
    private int[] id;

    public UnionFind(int n) {
        id = new int[n];

        for (int i = 0; i < n; i++)
            id[i] = i;
    }

    public int find(int p) {
        if (id[p] == p)
            return p;
        
        return id[p] = find(id[p]);
    }

    public void connect(int p, int q) {
        int pId = id[p];
        int qId = id[q];
        
        if (pId != qId)
            id[pId] = qId;
    }
}

public List<Connection> lowestCost(List<Connection> connections) {
    List<Connection> result = new ArrayList<>();

    Collections.sort(connections, new Comparator<Connection>() {
        @Override
        public int compare(Connection o1, Connection o2) {
            if (o1.cost != o2.cost)
                return o1.cost - o2.cost;
            else if (!o1.city1.equals(o2.city1))
                return o1.city1.compareTo(o2.city1);
            else
                return o1.city2.compareTo(o2.city2);
        }
    });

    Map<String, Integer> map = new HashMap<>();
    int index = 0;

    for (Connection c : connections) {
        if (!map.containsKey(c.city1))
            map.put(c.city1, index++);

        if (!map.containsKey(c.city2))
            map.put(c.city2, index++);
    }

    UnionFind uf = new UnionFind(index);

    for (Connection c : connections) {
        int idx1 = map.get(c.city1);
        int idx2 = map.get(c.city2);

        if (uf.find(idx1) != uf.find(idx2)) {
            result.add(c);
            uf.connect(idx1, idx2);
        }
    }

    return result.size() == index - 1 ? result : new ArrayList<>();
}