import java.util.HashSet;
/**
* Created by lxw, liwei4939@126.com on 2017/10/24.
*/
public class topKSum {
public static void main(String[] args){
int[] arr1 = {1,2,3,4,5};
int[] arr2 = {3,5,7,9,11};
int topK = 4;
int[] res = topK(arr1, arr2, topK);
for(int tmp:res)
System.out.print(tmp+" ");
}
public static class HeapNode{
public int row;
public int col;
public int value;
public HeapNode(int row, int col, int value){
this.row = row;
this.col = col;
this.value = value;
}
}
public static int[] topK(int[] arr1, int[] arr2, int topK){
if(arr1==null || arr2==null || topK <1)
return null;
topK = Math.min(topK, arr1.length * arr2.length);
HeapNode[] heap = new HeapNode[topK+1];
int heapSize =0;
int headR= arr1.length-1;
int headC= arr2.length-1;
int uR=-1;
int uC=-1;
int lR=-1;
int lC=-1;
heapInsert(heap, heapSize++, headR, headC, arr1[headR]+arr2[headC]);
HashSet<String> positionSet = new HashSet<String>();
int[] res = new int[topK];
int resindex=0;
while (resindex != topK){
HeapNode head = popHead(heap, heapSize--);
res[resindex++] = head.value;
headR = head.row;
headC = head.col;
uR = headR -1;
uC = headC;
if(headR != 0 && !isContains(uR, uC, positionSet)){
heapInsert(heap, heapSize++, uR, uC, arr1[uR]+arr2[uC]);
addPositionToSet(uR, uC, positionSet);
}
lR= headR;
lC = headC-1;
if(headC != 0 && !isContains(lR, lC, positionSet)){
heapInsert(heap, heapSize++, lR,lC, arr1[lR]+arr2[lC]);
addPositionToSet(lR, lC, positionSet);
}
}
return res;
}
public static HeapNode popHead(HeapNode[] heap, int heapSize){
HeapNode res = heap[0];
swap(heap, 0, heapSize-1);
heap[--heapSize] = null;
heapify(heap, 0, heapSize);
return res;
}
public static void heapify(HeapNode[] heap, int index, int heapSize){
int left = index*2 + 1;
int right = index*2 +2;
int largest = index;
while (left < heapSize){
if(heap[left].value > heap[index].value)
largest = left;
if(right<heapSize && heap[right].value > heap[largest].value)
largest = right;
if(index != largest){
swap(heap, largest, index);
} else {
break;
}
index = largest;
left = index * 2 +1;
right = index *2 +2;
}
}
public static void heapInsert(HeapNode[] heap, int index, int row, int col, int value){
heap[index] = new HeapNode(row, col, value);
int parent = (index-1)/2;
while (index != 0){
if(heap[index].value> heap[parent].value){
swap(heap, index, parent);
index = parent;
parent = (index-1)/2;
}else{
break;
}
}
}
public static void swap(HeapNode[] heap, int index1, int index2){
HeapNode tmp = heap[index1];
heap[index1] = heap[index2];
heap[index2]= tmp;
}
public static boolean isContains(int row, int col, HashSet<String> set){
return set.contains(String.valueOf(row + "_" + col));
}
public static void addPositionToSet(int row, int col, HashSet<String> set){
set.add(String.valueOf(row+ "_" + col));
}
}
两个有序数组相加和的topK问题
最新推荐文章于 2022-11-16 16:05:47 发布