前言

当前所有算法都使用测试用例运行过,但是不保证100%的测试用例,如果存在问题务必联系批评指正~

问题介绍

原问题
给定两个有序的数组,求两个有序数组之间两两数之和的topK。
如:arr1 = [1,2,3,4,6], arr2=[4,6,8,9], k=2
结果为:[14,15]

解决方案

原问题
建立一个堆,两个数组分别建立两个游标,则两两组合可以看作是矩阵,这里将行坐标row作为arr1的游标,col作为列arr2的游标
其中堆元素记录(row,col,arr1[row] + arr2[col])
初始化时将arr1[len-1] + arr2[len-1]放入堆中,后续每次弹出堆顶,并将堆顶的row和col拿出来,将当前坐标的左边(row,col-1),上边(row-1, col)放入到堆中重新进行堆上浮即可。

代码编写

java语言版本

原问题:
原问题:

  /**
     * 获取两个有序数组中的相加和topk
     * @param arr1
     * @param arr2
     * @return
     */
    public static int[] getTopKFromArrCp1(int[] arr1, int[] arr2, int k) {
        if (arr1 == null || arr2 == null
                || arr1.length == 0 || arr2.length == 0
                || arr1.length < k || arr2.length < k) {
            return null;
        }
        int row = arr1.length-1;
        int col = arr2.length-1;
        // 存储topk的容器
        RowColNode[] heap = new RowColNode[k+1];
        heap[0] = new RowColNode(row, col, arr1[row] + arr2[col]);
        int count = k;
        int[] res = new int[k];
        int heapSize = 1;
        while (k > 0) {
            // 弹出堆顶
            RowColNode head = heap[0];
            res[count-k] = head.getValue();
            CommonUtils.swapPlus(heap, 0, heapSize-1);
            heap[heapSize-1] = null;
            heapSize--;
            heapifyCp1(heap, heapSize);
            // 加入左边和上边的两个元素
            row = head.getRow();
            col = head.getCol();
            if (row - 1 >= 0) {
                heap[heapSize++] = new RowColNode(row-1, col, arr1[row-1] + arr2[col]);
                // 上浮
                heapInsertCp(heap, heapSize);
            }
            if (col - 1 >= 0) {
                heap[heapSize++] = new RowColNode(row, col-1, arr1[row] + arr2[col-1]);
                heapInsertCp(heap, heapSize);
            }
            k--;
        }
        return res;
    }

    private static void heapInsertCp(RowColNode[] heap, int heapSize) {
        int cur = heapSize-1;
        int parent = (cur-1)/2;
        while (cur != 0) {
            if (heap[cur].getValue() > heap[parent].getValue()) {
                CommonUtils.swapPlus(heap, cur, parent);
                cur = parent;
            }else {
                break;
            }
        }
    }

    private static void heapifyCp1(RowColNode[] heap, int heapSize) {
        int parent = 0;
        int left = (parent * 2) + 1;
        while (left < heapSize) {
            int maxIndex = heap[parent].getValue() > heap[left].getValue() ? parent : left;
            int right = left + 1;
            while (right < heapSize && heap[right].getValue() >heap[maxIndex].getValue()) {
                maxIndex = right;
            }
            if (maxIndex == parent) {
                break;
            }
            CommonUtils.swapPlus(heap, parent, maxIndex);
            parent = maxIndex;
            left = (parent * 2) + 1;
        }
    }


    public static void main(String[] args) {
        CommonUtils.printArr(getTopKFromArrCp1(new int[]{2,3,4,5,6}, new int[]{5,6,7,8,9}, 3));
    }


    public static class RowColNode {

        int row;
        int col;
        int value;

        public RowColNode(int row, int col, int value) {
            this.row = row;
            this.col = col;
            this.value = value;
        }

        public int getRow() {
            return row;
        }

        public void setRow(int row) {
            this.row = row;
        }

        public int getCol() {
            return col;
        }

        public void setCol(int col) {
            this.col = col;
        }

        public int getValue() {
            return value;
        }

        public void setValue(int value) {
            this.value = value;
        }
    }

c语言版本

正在学习中

c++语言版本

正在学习中

思考感悟

topK问题基本上99%都是往堆方面去思考,这道题也不例外,只是在于以坐标的方式入堆这个设计可以借鉴一下,特别是双游标甚至三游标时都可以通过矩阵或者张量(tensor)的形式去解决问题

写在最后

06-14 01:11