[자료구조] 힙 (Heap) - java

package com.algorithm2020;
import java.util.ArrayList;

enum HeapType {
    MIN,
    MAX
}

public class Heap {
    ArrayList<Integer> arr;
    int n; // 남아 있는 노드 개수
    HeapType type;

    // 생성자
    public Heap(ArrayList<Integer> arr, int firstN, HeapType type) {
        this.arr = arr;
        this.n = firstN;
        this.type = type;

        if (type.equals(HeapType.MIN)) {
            sortMinHeap(n);
        } else {
            sortMaxHeap(n);
        }
    }

    public int top() {
        return arr.get(1);
    }

    public void push(int insertValue) {
        // 제일 마지막에 insertValue를 넣는다.
        arr.add(insertValue);
        ++n;

        if (type.equals(HeapType.MIN)) {
            sortMinHeap(n);
        } else {
            sortMaxHeap(n);
        }
    }

    public int pop() {
        int rtn = 0;

        swap(1, n);

        // n-1 ~ 2 까지 sort!
        if (type.equals(HeapType.MIN)) {
            sortMinHeap(n-1);
        } else {
            sortMaxHeap(n-1);
        }

        rtn = arr.remove(n--);

        return rtn;
    }

    private void sortMinHeap(int lastIndex) {
        for (int i=lastIndex; i>1; i--) {
            if (arr.get(i/2) > arr.get(i)) swap(i / 2, i);
        }
    }

    private void sortMaxHeap(int lastIndex) {
        for (int i=lastIndex; i>1; i--) {
            if (arr.get(i/2) < arr.get(i)) swap(i/2, i);
        }
    }

    private void swap(int index1, int index2) {
        int tmp = arr.get(index1);
        arr.set(index1, arr.get(index2));
        arr.set(index2, tmp);
    }
}