Sample an index of a maximal number in an array, with a probability of 1/(number of maximal numbers)

Question

This is one of the recent interview question that I faced. Program to return the index of the maximum number in the array [ To Note : the array may or may not contain multiple copies of maximum number ] such that each index ( which contains the maximum numbers ) have the probability of 1/no of max numbers to be returned.

Examples:

• [-1 3 2 3 3], each of positions [1,3,4] have the probability 1/3 to be returned (the three 3s)
• [ 2 4 6 6 3 1 6 6 ], each of [2,3,6,7] have the probability of 1/4 to be returned (corresponding to the position of the 6s).

First, I gave O(n) time and O(n) space algorithm where I collect the set of max-indexes and then return a random number from the set. But he asked for a O(n) time and O(1) complexity program and then I came up with this.

int find_maxIndex(vector<int> a)
{
     max = a[0];
     max_index = 0;
     count = 0;

     for(i = 1 to a.size())
     {
         if(max < a[i])
         {
             max = a[i];
             count = 0;
         }
         if(max == a[i])
         {
              count++;
              if(rand < 1/count) //rand = a random number in the range of [0,1]
                  max_index = i;
         }
      }
      return max_index;
}

I gave him this solution. But my doubt is if this procedure would select one of the indexes of max numbers with equal probability. Hope I am clear.Is there any other method to do this ?

Answer 1

What you have is Reservoir sampling! There is another easy to understand solution, but requires two passes.

int find_maxIndex(vector<int> a){
    int count = 1;
    int maxElement = a[0];
    for(int i = 1; i < a.size(); i++){
        if(a[i] == maxElement){
            count ++;
        } else if(a[i] > maxElement){
            count = 1;
            maxElement = a[i];
        }
    }
    int occurrence = rand() % count + 1;
    int occur = 0;
    for(int i = 0; i < a.size(); i++){
        if(a[i] == maxElement){
            occur++;
            if(occur == occurrence) return i;
        }
    }
}

The algorithm is pretty simple, first find the number of times the max element occurs in the first pass. And choose a random occurrence and return the index of that occurrence. It takes two passes though, but very easy to understand.