TopK 主函数
| 44 | |
| 45 | // TopK 主函数 |
| 46 | void topKSelection(std::vector<int>& result, std::vector<int>& arr, int K, int maxIterations) { |
| 47 | int n = arr.size(); |
| 48 | int k = maxIterations; |
| 49 | result.clear(); |
| 50 | |
| 51 | |
| 52 | |
| 53 | while (k > 0) { |
| 54 | printf("K: %d\n", K); |
| 55 | printf("input size: %d\n", arr.size()); |
| 56 | // for(int i = 0; i < arr.size(); i++){ |
| 57 | // printf("%d ", arr[i]); |
| 58 | // } |
| 59 | // printf("\n"); |
| 60 | // 确定 pivots 数量 |
| 61 | int p = std::ceil(std::pow(n, 1.0 / k)); |
| 62 | p = std::min(p, n); // 防止 pivots 数量超过数组大小 |
| 63 | printf("p: %d\n", p); |
| 64 | |
| 65 | // 选择 pivots |
| 66 | std::vector<int> pivots = selectPivots(arr, p - 1); |
| 67 | // printf("pivots: "); |
| 68 | // for(int i = 0; i < pivots.size(); i++){ |
| 69 | // printf("%d ", pivots[i]); |
| 70 | // } |
| 71 | // printf("\n"); |
| 72 | |
| 73 | // 根据 pivots 对数组进行分区 |
| 74 | std::vector<std::vector<int>> blocks = partition(arr, pivots); |
| 75 | printf("blocks: "); |
| 76 | for(int i = 0; i < blocks.size(); i++){ |
| 77 | printf("["); |
| 78 | printf("%d", blocks[i].size()); |
| 79 | printf("] "); |
| 80 | } |
| 81 | printf("\n"); |
| 82 | |
| 83 | // 统计区块大小，找到包含 TopK 的区块 |
| 84 | int total = 0; |
| 85 | int targetBlock = -1; |
| 86 | arr.clear(); |
| 87 | for (int i = blocks.size() - 1; i >= 0; i--) { |
| 88 | if (total + blocks[i].size() >= K) { |
| 89 | arr.insert(arr.end(), blocks[i].begin(), blocks[i].end()); |
| 90 | break; |
| 91 | } |
| 92 | result.insert(result.end(), blocks[i].begin(), blocks[i].end()); |
| 93 | total += blocks[i].size(); |
| 94 | } |
| 95 | // printf("result: "); |
| 96 | // for(int i = 0; i < result.size(); i++){ |
| 97 | // printf("%d ", result[i]); |
| 98 | // } |
| 99 | // printf("\n"); |
| 100 | // 更新 K 和 arr |
| 101 | K -= total; |
| 102 | n = arr.size(); |
| 103 |