kth smallest element in an array

Question from lintcode

Given an unsorted array, find the kth smallest element.

Idea

Use any sort algorithm you are familiar with. I implemented a heap with an array. This array implementation is brought from Mark Allen Weiss 's Data Structures and Algorithm Analysis in C in my university class.

public class Solution {
    /*
     * @param k: An integer
     * @param nums: An integer array
     * @return: kth smallest element
     */
    public int kthSmallest(int k, int[] nums) {

        assert nums.length > 0;

        Heap heap = new Heap(nums.length);
        for(int n: nums) {
            heap.add(n);
        }

        assert heap.size() > 0 && k <= heap.size();

        int out = Integer.MIN_VALUE;
        for(int i = 0; i < k; i++) {
            out = heap.pop();
        }
        return out;
    }
    
    class Heap {
        private int[] hp;
        private int lastPos = 0;

        private int childPosLeft(int n) {
            return n * 2;
        }

        private int childPosRight(int n) {
            return childPosLeft(n) + 1;
        }

        private int parentPos(int n ) {
            return n/2;
        }

        private int root() {
            return 1;
        }

        private boolean hasChild(int pos) {
            return childPosLeft(pos) <= lastPos;
        }

        private  int smallerChildPos(int pos) {
            if (childPosLeft(pos) == lastPos) 
                return childPosLeft(pos); 
            if (hp[childPosLeft(pos)] > hp[childPosRight(pos)])
                return childPosRight(pos);
            return childPosLeft(pos);
        }

        public int size() {
            return lastPos;
        }

        public Heap(int size) {
            hp = new int[size + 1];
            hp[0] = Integer.MIN_VALUE;
        }

        public void add(int num) {
            lastPos++;
            // heap的精髓:只要你比我爸爸有錢(排序高, 這裏是越小越高),我叫我爸下來給你當兒子
            int i;
            for(i = lastPos; hp[parentPos(i)] > num; i = parentPos(i))
                hp[i] = hp[parentPos(i)];
            hp[i] = num;
        }

        public int pop() {
            int output = hp[root()];
            int i = root();
            int lastElement = hp[lastPos--];
            // 孫子上去頂層,往下走,見到爸爸們就讓他們向上爬,爬上去了空出來的位子你坐
            hp[i] = lastElement;
            while(hasChild(i)) {
                int child = smallerChildPos(i);
                if (hp[i] > hp[child]) {
                    hp[i] = hp[child];
                    hp[child] = lastElement;
                    i = child;
                } else {
                    break;
                }
            }
            return output;
        }
    }
}

你可能感兴趣的:(kth smallest element in an array)