定义一个基类,用基类的指针分别指向不同的子类。。。
#include <iostream>
#include <stdio.h>
#include <algorithm>
#include <vector>
using namespace std;
class ListNode {
public:
int val;
ListNode* next;
ListNode(int x) : val(x), next(NULL) {
}
};
class DataSource {
public:
ListNode* cur;
DataSource():cur(NULL) {}
virtual ListNode* getData() = 0;
};
class Disk: public DataSource {
public:
ListNode* getData() {
//maybe other implements
cur = NULL;
return NULL;
}
};
class Mem: public DataSource {
public:
ListNode* getData() {
// maybe other implements
cur = NULL;
return NULL;
}
};
class cmp {
public:
bool operator()(DataSource* a, DataSource* b) {
return a->cur->val > b->cur->val;
}
};
class Solution {
public:
ListNode *mergeKLists(vector<DataSource *> &lists) {
// Note: The Solution object is instantiated only once and is reused by each test case.
ListNode *head = NULL, *cur = NULL;
vector<DataSource*>::iterator it = lists.begin();
while (it != lists.end()) {
if ((*it)->getData() == NULL)
it = lists.erase(it);
else
it++;
}
if (lists.size() < 1)
return NULL;
make_heap(lists.begin(), lists.end(), cmp());
while (lists.size() > 0) {
if (head == NULL)
head = cur = lists[0]->cur;
else {
cur->next = lists[0]->cur;
cur = cur->next;
}
pop_heap(lists.begin(), lists.end(), cmp());
if (lists[lists.size() - 1]->getData() != NULL)
push_heap(lists.begin(), lists.end(), cmp());
else
lists.pop_back();
}
return head;
}
};
int main() {
DataSource* d1 = new Disk();
DataSource* d2 = new Mem();
vector<DataSource*> lists;
lists.push_back(d1);
lists.push_back(d2);
Solution s;
ListNode* res = s.mergeKLists(lists);
return 0;
}