FB面经Prepare: Merge K sorted Array

Merge K sorted Array

跟Merge K sorted lists不同在于，从PQ里poll出来以后不知道下一个需要被加入PQ的是哪一个

所以需要写一个wrapper class

 package fbPractise;

 import java.util.*;

 public class MergeKLists {
     static class Pair {
         int listIndex;
         int idInList;
         int value;
         public Pair(int l, int id, int val) {
             this.listIndex = l;
             this.idInList = id;
             this.value = val;
         }
     }

     public static List<Integer> merge(List<List<Integer>> lists) {
         List<Integer> res = new ArrayList<Integer>();

         Comparator<Pair> comp = new Comparator<Pair>() {
             public int compare(Pair p1, Pair p2) {
                 return p1.value - p2.value;
             }
         };
         PriorityQueue<Pair> pq = new PriorityQueue<Pair>(1, comp);
         for (int i=0; i<lists.size(); i++) {
             Pair p = new Pair(i, 0, lists.get(i).get(0));
             pq.offer(p);
         }

         while (!pq.isEmpty()) {
             Pair pa = pq.poll();
             int index = pa.listIndex;
             int id = pa.idInList;
             if (id < lists.get(index).size()-1) {
                 pq.offer(new Pair(index, id+1, lists.get(index).get(id+1)));
             }
             res.add(pa.value);
         }

         return res;
     }

     /**
      * @param args
      */
     public static void main(String[] args) {
         // TODO Auto-generated method stub
         List<Integer> l1 = Arrays.asList(1,2,2,3,6);
         List<Integer> l2 = Arrays.asList(1,4,5,7,8,9);
         List<Integer> l3 = Arrays.asList(3,3,3,5,10);
         List<List<Integer>> lists = new ArrayList<List<Integer>>();
         lists.add(new ArrayList<Integer>(l1));
         lists.add(new ArrayList<Integer>(l2));
         lists.add(new ArrayList<Integer>(l3));
         List<Integer> res = merge(lists);
         System.out.println(res);
     }

 }