leetcode 378 手写堆和堆的应用 附测试代码

非最佳方案，仅用于测试手撕堆的代码

class Heap{
    // big first heap;
    public int[] data = new int[100000];
    public int total = 0;

    public void insert(int x) {
        total++;
        data[total] = x;
        int cur = total;
        while (cur > 1) {
            int father = cur >> 1;
            if (data[father] < data[cur]) {
                int c = data[father];
                data[father] = data[cur];
                data[cur] = c;
                cur = father;
            } else {
                break;
            }
        }
    }

    public int top() {
        return data[1];
    }

    public void pop() {
        data[1] = data[total];
        total--;
        int cur = 1;
        while (true) {
            int son1 = cur * 2;
            int son2 = cur * 2 + 1;
            if (son1 > total) {
                break;
            } else if (son2 > total) {
                if (data[cur] < data[son1]) {
                    int c = data[cur];
                    data[cur] = data[son1];
                    data[son1] = c;
                }
                cur = son1;
            } else {
                int curSon = 0;
                if (data[son1] > data[son2]) {
                    curSon = son1;
                } else {
                    curSon = son2;
                }
                if (data[cur] < data[curSon]) {
                    int c = data[cur];
                    data[cur] = data[curSon];
                    data[curSon] = c;
                }
                cur = curSon;
            }
        }
    }

}


class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        Heap heap = new Heap();
        int n = matrix.length;
        int num = 1;
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < n; j++) {
//                System.out.print(i);
//                System.out.print(" ");
//                System.out.println(j);
                if (num <= k) {
                    heap.insert(matrix[i][j]);
                    num++;
                } else {
                    int heapTop = heap.top();
                    if (matrix[i][j] < heapTop) {
                        heap.pop();
                        heap.insert(matrix[i][j]);
                    }
                }
            }
        }
        return heap.top();
    }
}

附测试代码

import org.junit.Test;

import java.util.List;

import static org.junit.Assert.*;

public class SolutionTest {

    @Test
    public void test1() {
        int[][] matrix = new int[][]{{1, 5, 9}, {10, 11, 13}, {12, 13, 15}};
        int k = 8;
        Solution solution = new Solution();
        assertEquals(13, solution.kthSmallest(matrix, k));

        matrix = new int[][]{{-5}};
        k = 1;
        assertEquals(-5, solution.kthSmallest(matrix, k));

        matrix = new int[][]{{1, 2}, {1, 3}};
        k = 2;
        assertEquals(1, solution.kthSmallest(matrix, k));


    }

    @Test
    public void test2() {
        Solution solution = new Solution();
        int[][] matrix = new int[][]{{4,7,12},{4,9,17},{4,12,17}};
        int k = 3;
        assertEquals(4, solution.kthSmallest(matrix, k));
    }

    @Test
    public void test3() {
        Solution solution = new Solution();
        int[][] matrix = new int[][]{{1,4,7,11,15},{2,5,8,12,19},{3,6,9,16,22},{10,13,14,17,24},{18,21,23,26,30}};
        int k = 5;
        assertEquals(5, solution.kthSmallest(matrix, k));
    }


}