[LintCode] Load Balancer

Implement a load balancer for web servers. It provide the following functionality:

1. Add a new server to the cluster => add(server_id).
2. Remove a bad server from the cluster => remove(server_id).
3. Pick a server in the cluster randomly with equal probability => pick().

Example

At beginning, the cluster is empty => {}.

add(1)
add(2)
add(3)
pick()
>> 1         // the return value is random, it can be either 1, 2, or 3.
pick()
>> 2
pick()
>> 1
pick()
>> 3
remove(1)
pick()
>> 2
pick()
>> 3
pick()
>> 3

Analysis
To implement add/remove server id in O(1), we can only use hash map. However, it is not easy to 
implement pick randomly. 

Using array will be easy to implement pick random method. But using array alone does not give us O(1) remove operation.
Instead, we get a O(n) remove.

The solution is to use both a hash map and an arraylist to get O(1) add/remove/pick operations. 

The arraylist stores all server ids and the hash map maps server id to its index in the arraylist. 

Add:  add a new server id to the end of the arraylist; add a new key value pair for this newly added server id in the hashmap.

Remove:  Get the arraylist index of the server to be removed from the hash map, then remove the key value pair.

　　　　  Swap the server id at the end of the arraylist with the server id that is to be removed.

　　　　  Update the arraylist index of the swapped server id in the hash map. 

　　　　  Remove the server index at the end of the arraylist.

Pick:  use rand to randomly pick a arraylist index.

 

The key points of the above solution are:

1. To achieve O(1) deletion in the arraylist,

swap the element that needs to be removed with the last element, then delete the last element.

2. To achieve O(1) lookup of which server id to remove in the arraylist,

store a server id to arraylist index mapping in the hashmap.

 

public class LoadBalancer {
    private HashMap<Integer, Integer> map;
    private ArrayList<Integer> list;
    private Random rand;
    public LoadBalancer() {
        this.map = new HashMap<Integer, Integer>();
        this.list = new ArrayList<Integer>();
        this.rand = new Random();
    }

    // @param server_id add a new server to the cluster 
    // @return void
    public void add(int server_id) {
        if(!map.containsKey(server_id)){
            list.add(server_id);
            map.put(server_id, list.size() - 1);    
        }
    }

    // @param server_id server_id remove a bad server from the cluster
    // @return void
    public void remove(int server_id) {
        if(map.containsKey(server_id)){
            int removeIdx = map.get(server_id);
            map.remove(server_id);
            int affectedServerId = list.get(list.size() - 1);
            map.put(affectedServerId, removeIdx);
            list.set(removeIdx, affectedServerId);
            list.remove(list.size() - 1);
        }
    }

    // @return pick a server in the cluster randomly with equal probability
    public int pick() {
        if(list.size() == 0){
            return -1;
        }
        int rand_idx = rand.nextInt(list.size());
        return list.get(rand_idx);
    } 
}