C++ 多态 从不同数据源获取数据 多路归并

定义一个基类，用基类的指针分别指向不同的子类。。。

#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;
}