题解 | #最小的K个数#

最小的K个数

http://www.nowcoder.com/practice/6a296eb82cf844ca8539b57c23e6e9bf

注意partition方法参数是lo hi时,下面的坐标运算也要是lo hi相关而不是input.size()

堆和partition partition一开始没想到

堆可以是O(n)创建整个的堆然后O(klogn)获取最小的k个元素 也可以直接维护一个size为k的堆,O(nlogk)?


#define Parent(i) ((i - 1) >> 1)
#define LChild(i) (1 + (i << 1))
#define RChild(i) (2 + (i << 1))
#define LastInternal(n) Parent(n - 1)
class Solution {
public:
    void down(vector<int>& input, int idx){
        int val = input[idx];
        while(idx <= LastInternal(input.size())){
            int l = 1001, r =1001;
            int l_idx = LChild(idx);
            int r_idx = RChild(idx);
            if(l_idx < input.size()){
                l = input[l_idx];
            }
            if(r_idx < input.size()){
                r = input[r_idx];
            }
            if(l <= r){
                if(l < val){
                    input[idx] = l;
                    idx = l_idx;
                }
                else{
                    break;
                }
            }else{
                if(r < val){
                    input[idx] = r;
                    idx = r_idx;
                }else{
                    break;
                }
            }
        }
        input[idx] = val;
    }
    void up(vector<int>& input, int idx){
        int val = input[idx];
        while(idx >= 0){
            int p_idx = Parent(idx);
            int p_val = input[p_idx];
            if(p_val > val){
                input[idx] = p_val;
                idx = p_idx;
            }
            else{
                break;
            }
        }
        input[idx] = val;
    }
    void init_heap(vector<int>& input){
        for(int i = LastInternal(input.size()); i >= 0; i--){
            down(input, i);
        }
    }
    int del(vector<int>& input){
        int val = input[0];
        input[0] = input[input.size() - 1];
        input.pop_back();
        down(input, 0);
        return val;
        
    }
    void swap(int& a, int& b){
        a = a ^ b;
        b = a ^ b;
        a = a ^ b;
    }
    int partition(vector<int>& input, int lo, int hi){
        int mi = (lo + hi) >> 1;
        int val = input[mi];
        int pivot = lo;
        input[mi] = input[hi - 1];
        for(int i = lo ; i < hi - 1 ; i++){
            if(input[i] <= val){
                if(pivot != i){
                    swap(input[pivot], input[i]);
                }
                pivot++;
            }
        }
        input[hi - 1] = input[pivot];
        input[pivot] = val;
        return pivot;
    }
    vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
        vector<int> res;
        if(input.size() == 0 or k == 0){
            return res;
        }
        if(k >= input.size())
            return input;
//         init_heap(input);
//         for(int i = 0 ; i < k ; i++){
//             res.push_back(del(input));
//         }
        int mi = input.size() >> 1;
        int lo = 0;
        int hi = input.size();
        while(1){
            mi = partition(input, lo, hi);
            if(mi == k or mi == k - 1){
                return vector<int>(input.begin(), input.begin() + k);
            }
            else if(mi < k - 1){
                lo = mi + 1; 
            }
            else{
                hi = mi;
            }
        }
        return res;
        
    }
};
全部评论

相关推荐

头像
不愿透露姓名的神秘牛友
04-08 00:50
点赞 评论 收藏
转发
点赞 收藏 评论
分享
牛客网
牛客企业服务