O(log n) pick solution with binary search


  • 2
    E

    I'm surprised that O(n) pick time solutions pass the OJ. I thought that the goal of this exercise is to design data structure that allows faster than O(n) pick. Sharing my solution:

    class Solution {
    public:
        // Solution: 
        //  * initialization: keep array of (val, index) pairs. Sort the array by val.
        //  * pick: find lower and upper bound of target in the array; pick random element from the range
        
        struct data
        {
            int val;
            int ind;
        };
        
        std::default_random_engine generator;
        
        vector<data> counters;
        Solution(vector<int> nums) {
            srand (time(NULL));
            counters.resize(nums.size());
            for(int i=0; i < nums.size(); ++i)
                counters[i] = data{nums[i], i};
            sort(counters.begin(), counters.end(), [](const data &l, const data &r){return l.val < r.val;});
        }
        
        int pick(int target) {
            int lo = 0;
            int hi = counters.size()-1;
            while(hi>lo+1)
            {
                int mid = (hi + lo)/2;
                if(counters[mid].val < target)
                    lo=mid+1;
                else
                    hi=mid;
            }
            if(counters[lo].val != target)
                ++lo;
            
            int low_bound = lo;
            
            lo = 0, hi = counters.size()-1;
            while(hi>lo+1)
            {
                int mid = (hi + lo)/2;
                if(counters[mid].val > target)
                    hi=mid-1;
                else
                    lo=mid;
            }        
            if(counters[hi].val != target)
                --hi;
            int up_bound = hi;
            //std::uniform_int_distribution<int> distribution(low_bound,up_bound);
            int range = up_bound - low_bound;
            int ind = low_bound + (range != 0 ? rand() % (range+1) : 0 );// distribution(generator);
            
            return counters[ind].ind;
        }
    };
    

  • 0
    M

    @EgorAYusov
    I had the same impression -- namely that the pick() operation is the one that should be optimized. My code roughly mirrors yours, except for relying on equal_range() for the binary search. I tried both ways, this way and with the reservoir technique, and they produced similar run times :

        class Solution {
        public:
            Solution(const vector<int>& nums)
            {
                auto N = static_cast<int>(nums.size());
    
                items.resize(N);
                for (int i = 0; i < N; ++i) {
                    items[i] = { nums[i], i };
                }
    
                sort(begin(items), end(items));
            }
    
            int pick(int target)
            {
                auto range = equal_range(begin(items), end(items), item({ target, 0 }));
                auto index = distance(begin(items), range.first);
    
                auto count = distance(range.first, range.second);
                if (1 < count) {
                    index += uniform_int_distribution<int>(0, count - 1)(raneng);
                }
    
                return items[index].idx;
            }
    
    
            struct item {
                int val; 
                int idx;
    
                bool operator<(const item& rhs) const
                { return val < rhs.val; }
            };
    
            default_random_engine raneng;
            vector<item> items;
        };
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.