-
Notifications
You must be signed in to change notification settings - Fork 0
/
findTopKthInMatrix.cpp
81 lines (76 loc) · 1.77 KB
/
findTopKthInMatrix.cpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
#include <iostream>
#include <cstdio>
#include <cstdlib>
#include <queue>
#include <vector>
#include <climits>
using namespace std;
struct Node {
int val;
int x;
int y;
Node(int val, int x, int y) : val(val), x(x), y(y) {}
};
struct Comparator {
bool operator()(Node first, Node second) {
first.val > second.val; // min heap
}
};
vector<int> FirstKElement(vector<vector<int> > A, int k) {
priority_queue<Node, vector<Node>, Comparator> minHeap;
Node head (A[0][0], 0, 0);
vector<int> result;
minHeap.push(head);
int n = A.size();
for (int i = 0; i < k; i++) {
if (!minHeap.empty()) {
Node node = minHeap.top();
minHeap.pop();
while (!minHeap.empty() && minHeap.top().val == node.val) minHeap.pop();
if (node.x == n-1 && node.y == n-1) {
continue;
} else if (node.x == n-1) {
Node left(A[node.x][node.y+1], node.x, node.y+1);
minHeap.push(left);
} else if (node.y == n-1) {
Node down(A[node.x+1][node.y], node.x+1, node.y);
minHeap.push(down);
} else {
Node left(A[node.x][node.y+1], node.x, node.y+1);
minHeap.push(left);
Node down(A[node.x+1][node.y], node.x+1, node.y);
minHeap.push(down);
}
result.push_back(node.val);
}
}
return result;
}
int main() {
vector<int> vec(10, 0);
for (int i = 0; i < 10; i++) {
vec[i] = i;
}
vector<vector<int> > A;
for (int i = 0; i < 10; i++) {
vector<int> vec1 = vec;
for (int j = 0; j < 10; j++) {
vec1[j] += j;
}
A.push_back(vec1);
}
for (int i = 0; i < 10; i++) {
for (int j = 0; j < 10; j++) {
std::cout << A[i][j] << " ";
}
std::cout << std::endl;
}
std::cout << std::endl;
int k = 10;
vector<int> result = FirstKElement(A, k);
for (int i = 0; i < result.size(); i++) {
std::cout << result[i] << " ";
}
std::cout << std::endl;
return 0;
}