FB面经Prepare: Merge K sorted Array

Merge K sorted Array

1 package fbPractise;
2
3 import java.util.*;
4
5 public class MergeKLists {
6     static class Pair {
7         int listIndex;
8         int idInList;
9         int value;
10         public Pair(int l, int id, int val) {
11             this.listIndex = l;
12             this.idInList = id;
13             this.value = val;
14         }
15     }
16
17     public static List<Integer> merge(List<List<Integer>> lists) {
18         List<Integer> res = new ArrayList<Integer>();
19
20         Comparator<Pair> comp = new Comparator<Pair>() {
21             public int compare(Pair p1, Pair p2) {
22                 return p1.value - p2.value;
23             }
24         };
25         PriorityQueue<Pair> pq = new PriorityQueue<Pair>(1, comp);
26         for (int i=0; i<lists.size(); i++) {
27             Pair p = new Pair(i, 0, lists.get(i).get(0));
28             pq.offer(p);
29         }
30
31         while (!pq.isEmpty()) {
32             Pair pa = pq.poll();
33             int index = pa.listIndex;
34             int id = pa.idInList;
35             if (id < lists.get(index).size()-1) {
36                 pq.offer(new Pair(index, id+1, lists.get(index).get(id+1)));
37             }
38             res.add(pa.value);
39         }
40
41         return res;
42     }
43
44     /**
45      * @param args
46      */
47     public static void main(String[] args) {
48         // TODO Auto-generated method stub
49         List<Integer> l1 = Arrays.asList(1,2,2,3,6);
50         List<Integer> l2 = Arrays.asList(1,4,5,7,8,9);
51         List<Integer> l3 = Arrays.asList(3,3,3,5,10);
52         List<List<Integer>> lists = new ArrayList<List<Integer>>();
53         lists.add(new ArrayList<Integer>(l1));
54         lists.add(new ArrayList<Integer>(l2));
55         lists.add(new ArrayList<Integer>(l3));
56         List<Integer> res = merge(lists);
57         System.out.println(res);
58     }
59
60 }