无序数组前K个最大数

无序数组前K个最大数

思路：使用前K个元素建立一个K个元素的小顶堆，然后遍历数组后面的内容替换堆顶并保持小顶堆。

/* 
参考链接：https://github.com/yxjxx/SortAlgorithmsEx/blob/master/SortAlgs/yj_quick_sort.cpp
https://blog.csdn.net/BRO_BMW/article/details/105985158
*/

@implementation ViewController

- (void)viewDidLoad {
    [super viewDidLoad];   

    NSArray *arr = @[@22, @4, @7, @1, @2, @3, @5, @3, @6, @3, @2, @18, @7];
    NSArray *n = [self findK:arr k:5];
    NSLog(@"%@", n);//(6,7,7,18,22)
}

- (NSArray<NSNumber *> *)findK:(NSArray *)arr k:(int)k{
    NSMutableArray *heap = [self quickSortFirstK:arr k:k];
    for (int i = k; i<arr.count; i++) {
        if ([arr[i] intValue] > [heap[0] intValue]) {
            heap[0] = arr[i];
            [self heapify:heap idx:0];
        }
    }
    return heap;
}

//节点 i 的左子节点下标2i+1， 右子节点下标2i+2
- (void)heapify:(NSMutableArray *)arrM idx:(int)idx{
    if (idx*2 + 1 > arrM.count-1) {
        return;
    }
    if ([arrM[idx] intValue] > [arrM[2*idx+1] intValue]) {
        [self swapIndexA:idx indexB:2*idx+1 arrM:arrM];
        [self heapify:arrM idx:2*idx+1];
        if ((2*idx+2) <= arrM.count-1 && [arrM[idx] intValue] > [arrM[2*idx+2] intValue]) {
            [self swapIndexA:idx indexB:2*idx+2 arrM:arrM];
            [self heapify:arrM idx:2*idx+2];
        }
    } else if ((2*idx+2) <= arrM.count-1 && [arrM[idx] intValue]> [arrM[2*idx+2] intValue]) {
        [self swapIndexA:idx indexB:2*idx+2 arrM:arrM];
        [self heapify:arrM idx:2*idx+2];
    }
    return;
}

- (void)swapIndexA:(int)a indexB:(int)b arrM:(NSMutableArray *)arrM {
    NSNumber *temp = arrM[a];
    arrM[a] = arrM[b];
    arrM[b] = temp;
}

- (NSMutableArray *)quickSortFirstK:(NSArray *)arr k:(int)k {
    NSArray *subArr = [arr subarrayWithRange:NSMakeRange(0, k)];
    NSMutableArray *arrM = [subArr mutableCopy];
    [self quickSort:arrM start:0 end:k-1];
    return arrM;
}

- (void)quickSort:(NSMutableArray *)arrM start:(int)start end:(int)end {
    if (start >= end) {
        return;
    }//递归结束条件
    int pivot = [arrM[start] intValue];
    int left = start;
    int right = end;
    while (left < right) {
        while (left < right && [arrM[right] intValue] > pivot) {
            right--;
        }
        if (left < right) {
            arrM[left] = arrM[right];
            left++;
        }
        
        while (left < right && [arrM[left] intValue] < pivot) {
            left++;
        }
        
        if (left < right) {
            arrM[right] = arrM[left];
            right--;
        }
    }
    arrM[left] = @(pivot);
    
    [self quickSort:arrM start:start end:left-1];
    [self quickSort:arrM start:left+1 end:end];
}

@end