改变已知排序的key，依然保持大根堆或者小根堆

Posted 2021-02-20 laydown

import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashMap;

/**
 * 改变已知排序的key，依然保持大根堆或者小根堆
 */
public class ChangeSortKeyHeap {

    public static class MyHeap<T> {

        public ArrayList<T> heap;

        public HashMap<T, Integer> indexMap;

        public int heapSize;

        public Comparator<T> comp;

        public MyHeap(Comparator comp) {
            this.comp = comp;
            this.heap = new ArrayList<>();
            indexMap = new HashMap<>();
            heapSize = 0;
        }

        public void push(T value) {
            if (contains(value)) {
                heap.set(indexMap.get(value), value);
                return;
            }
            heap.add(value);
            indexMap.put(value, heapSize);
            heapInsert(heapSize++);
        }

        public T pop() {
            if (isEmpty()) {
                System.out.println("the heap is empty");
                return null;
            }
            T res = heap.get(0);
            int end = heapSize - 1;
            swap(0, end);
            heap.remove(end);
            indexMap.remove(res);
            heapify(0, --heapSize);
            return res;
        }

        /**
         * 修改对象的key，并重新使有序
         *
         * @param value 对象
         */
        public void resign(T value) {
            if (!contains(value)) {
                System.out.println("the value is not in the heap");
                return;
            }
            int valueIndex = indexMap.get(value);
            heapInsert(valueIndex);
            heapify(valueIndex, heapSize);
        }

        /**
         * 修改了一个元素之后，保持原逻辑排序
         *
         * @param index    修改的元素位置
         * @param heapSize 堆大小
         */
        public void heapify(int index, int heapSize) {
            int left = (index << 1) + 1;
            while (left < heapSize) {
                int largest = left + 1 < heapSize && comp.compare(heap.get(left + 1), heap.get(left)) < 0 ? left + 1 : left;
                if (comp.compare(heap.get(index), heap.get(largest)) <= 0) {
                    break;
                }
                swap(index, largest);
                index = largest;
                left = (index << 1) + 1;
            }
        }

        /**
         * 往堆中添加一个元素，同时保持原逻辑排序
         *
         * @param index 新添加的元素的位置
         */
        public void heapInsert(int index) {
            int pIndex = 0;
            // 此处 ((index - 1) / 2) 和 ((index - 1) >> 1) 计算结果是不一样的，当index=0，前者是0，后者是-1
            while (comp.compare(heap.get(index), heap.get(pIndex = ((index - 1) / 2))) < 0) {
                swap(index, pIndex);
                index = pIndex;
            }
        }

        /**
         * 交换两个元素的位置
         *
         * @param i 位置
         * @param j 位置
         */
        private void swap(int i, int j) {
            // 同一个位置交换无意义
            if (i == j) {
                return;
            }
            // 交换
            T o1 = heap.get(i);
            T o2 = heap.get(j);
            heap.set(i, o2);
            heap.set(j, o1);
            indexMap.put(o1, j);
            indexMap.put(o2, i);
        }

        public int size() {
            return heapSize;
        }

        public boolean isEmpty() {
            return heapSize == 0;
        }

        public boolean contains(T value) {
            if (isEmpty()) {
                return false;
            }
            return indexMap.containsKey(value);
        }

    }

}

/* 如有意见或建议，欢迎评论区留言；如发现代码有误，欢迎批评指正 */