使用加強堆結構解決topK問題

Grey Zeng發表於2022-04-10

作者:Grey

原文地址: 使用加強堆結構解決topK問題

題目描述

LintCode 550 · Top K Frequent Words II

思路

由於要統計每個字串的次數,以及字典序,所以,我們需要把使用者每次add的字串封裝成一個物件,這個物件中包括了這個字串和這個字串出現的次數。

假設我們封裝的物件如下:

public class Word {
    public String value; // 對應的字串
    public int times; // 對應的字串出現的次數

    public Word(String v, int t) {
        value = v;
        times = t;
    }
}

topk的要求是: 出現次數多的排前面,如果次數一樣,字典序小的排前面

很容易想到用有序表+比較器來做。

比較器的規則定義成和topk的要求一樣,然後把元素元素加入使用比較器的有序表中,如果要返回topk,直接從這個有序表彈出返回給使用者即可。比較器的定義如下:

public class TopKComparator implements Comparator<Word> {
    @Override
    public int compare(Word o1, Word o2) {
        // 次數大的排前面,次數一樣字典序在小的排前面
        return o1.times == o2.times ? o1.value.compareTo(o2.value) : (o2.times - o1.times);
    }
}

有序表配置這個比較器即可

TreeSet<Word>  topK = new TreeSet<>(new TopKComparator());

所以topk()方法很簡單,只需要從有序表裡面把元素拿出來返回給使用者即可

public List<String> topk() {
      List<String> result = new ArrayList<>();
      for (Word word : topK) {
          result.add(word.value);
      }
      return result;
}

時間複雜度 O(K)

以上步驟不復雜,接下來是add的邏輯,add的每次操作都有可能對前面我們設定的topK有序表造成影響,

所以在每次add操作的時候需要有一個機制可以告訴topK這個有序表,需要淘汰什麼元素,需要新加哪個元素,讓topK這個有序表時時刻刻只存topk個元素,

這樣就可以確保topK()方法比較單純,時間複雜度保持在O(K)

所以接下來的問題是:如何告訴topK這個有序表,需要淘汰什麼元素,需要新加哪個元素?

我們可以通過堆來維持一個門檻,堆頂元素表示最先要淘汰的元素,所以堆中的比較策略定為:

次數從小到大,字典序從大到小,這樣,堆頂元素永遠是:次數相對更少或者字典序相對更大的那個元素。所以如果某個時刻要淘汰一個元素,從堆頂拿出來,然後再到topK這個有序表中查詢是否有這個元素,有的話就從topK這個有序表中刪除這個元素即可。

private class ThresholdComparator implements Comparator<Word> {

    @Override
    public int compare(Word o1, Word o2) {
        // 設定堆門檻,堆頂元素最先被淘汰
        return o1.times == o2.times ? o2.value.compareTo(o2.value) : (o1.times - o2.times);
    }
}

如果使用Java自帶的PriorityQueue做這個堆,無法實現動態調整堆的功能,因為我們需要把次數增加的字串(Word)在堆上動態調整,自帶的PriorityQueue無法實現這個功能,PriorityQueue只能支援每次新增或者刪除一個節點的時候,動態調整堆(
O(logN),但是如果堆中的節點變化了,PriorityQueue無法自動調整成堆結構,所以我們需要實現一個增強堆,用於節點變化的時候可以動態調整堆結構(保持O(logN)複雜度)。

加強堆的核心是增加了一個雜湊表,

private Map<Word, Integer> indexMap;

用於存放每個節點所在堆上的位置,在節點變化的時候,可以通過雜湊表查出這個節點所在的位置,然後從所在位置進行heapify/heapInsert操作,且這兩個操作只會走一個,
這樣就動態調整好了這個堆結構,以下resign方法就是完成這個工作

public void resign(Word word) {
    int i = indexMap.get(word);
    heapify(i);
    heapInsert(i);
}

除了這個resign方法,自定義堆中的其他方法和常規的堆沒有區別,在每次進行heapify和heapInsert操作的時候,如果涉及到交換兩個元素,需要將indexMap中的兩個元素的位置也互換

private void swap(int i, int j) {
 if (i != j) {
  indexMap.put(words[i], j);
  indexMap.put(words[j], i);
  Word tmp = words[i];
  words[i] = words[j];
  words[j] = tmp;
 }
}

由於自定義堆和有序表topk只存top k個資料,所以TopK結構中還需要一個雜湊表來記錄所有的字串出現與否:

private Map<String, Word> map;

自此,TopK結構中的add方法需要的前置條件已經具備,整個add方法的流程如下:

image

關於複雜度,add方法,時間複雜度O(log K), topk方法,時間複雜度O(K)

完整程式碼

class TopK {
    private TreeSet<Word> topK;
    private Heap heap;
    private Map<String, Word> map;
    private int k;

    public TopK(int k) {
        this.k = k;
        topK = new TreeSet<>(new TopKComparator());
        heap = new Heap(k, new ThresholdComparator());
        map = new HashMap<>();
    }

    public void add(String str) {
        if (k == 0) {
            return;
        }
        Word word = map.get(str);
        if (word == null) {
            // 新增元素
            word = new Word(str, 1);
            // 是否到達門檻可以替換堆中元素
            if (heap.isReachThreshold(word)) {
                if (heap.isFull()) {
                    Word toBeRemoved = heap.poll();
                    topK.remove(toBeRemoved);
                }
                heap.add(word);
                topK.add(word);
            }
        } else {
            if (heap.contains(word)) {
                topK.remove(word);
                word.times++;
                topK.add(word);
                heap.resign(word);
            } else {
                word.times++;
                if (heap.isReachThreshold(word)) {
                    if (heap.isFull()) {
                        Word toBeRemoved = heap.poll();
                        topK.remove(toBeRemoved);
                    }
                    heap.add(word);
                    topK.add(word);
                }
            }
        }
        map.put(str, word);
    }

    public List<String> topk() {
        if (k == 0) {
            return new ArrayList<>();
        }
        List<String> result = new ArrayList<>();
        for (Word word : topK) {
            result.add(word.value);
        }
        return result;
    }

    private class Word {
        public String value;
        public int times;

        public Word(String v, int t) {
            value = v;
            times = t;
        }
    }

    private class TopKComparator implements Comparator<Word> {
        @Override
        public int compare(Word o1, Word o2) {
            // 次數大的排前面,次數一樣字典序在小的排前面
            return o1.times == o2.times ? o1.value.compareTo(o2.value) : (o2.times - o1.times);
        }
    }

    private class ThresholdComparator implements Comparator<Word> {

        @Override
        public int compare(Word o1, Word o2) {
            // 設定堆門檻,堆頂元素最先被淘汰
            return o1.times == o2.times ? o2.value.compareTo(o1.value) : (o1.times - o2.times);
        }
    }

    private class Heap {
        private Word[] words;
        private Comparator<Word> comparator;
        private Map<Word, Integer> indexMap;

        public Heap(int k, Comparator<Word> comparator) {
            words = new Word[k];
            indexMap = new HashMap<>();
            this.comparator = comparator;
        }

        public boolean isEmpty() {
            return indexMap.isEmpty();
        }

        public boolean isFull() {
            return indexMap.size() == words.length;
        }

        public boolean isReachThreshold(Word word) {
            if (isEmpty() || indexMap.size() < words.length) {
                return true;
            } else {
                if (comparator.compare(words[0], word) < 0) {
                    return true;
                }
                return false;
            }
        }

        public void add(Word word) {
            int size = indexMap.size();
            words[size] = word;
            indexMap.put(word, size);
            heapInsert(size);

        }

        private void heapify(int i) {
            int size = indexMap.size();
            int leftChildIndex = 2 * i + 1;
            while (leftChildIndex < size) {
                Word weakest = leftChildIndex + 1 < size
                        ? (comparator.compare(words[leftChildIndex], words[leftChildIndex + 1]) < 0
                        ? words[leftChildIndex]
                        : words[leftChildIndex + 1])
                        : words[leftChildIndex];
                if (comparator.compare(words[i], weakest) < 0) {
                    break;
                }
                int weakestIndex = weakest == words[leftChildIndex] ? leftChildIndex : leftChildIndex + 1;
                swap(weakestIndex, i);
                i = weakestIndex;
                leftChildIndex = 2 * i + 1;
            }
        }

        public void resign(Word word) {
            int i = indexMap.get(word);
            heapify(i);
            heapInsert(i);
        }

        private void heapInsert(int i) {
            while (comparator.compare(words[i], words[(i - 1) / 2]) < 0) {
                swap(i, (i - 1) / 2);
                i = (i - 1) / 2;
            }
        }

        public boolean contains(Word word) {
            return indexMap.containsKey(word);
        }

        public Word poll() {
            Word result = words[0];
            swap(0, indexMap.size() - 1);
            indexMap.remove(result);
            heapify(0);
            return result;
        }

        private void swap(int i, int j) {
            if (i != j) {
                indexMap.put(words[i], j);
                indexMap.put(words[j], i);
                Word tmp = words[i];
                words[i] = words[j];
                words[j] = tmp;
            }
        }
    }
}

更多

演算法和資料結構筆記

參考資料

演算法和資料結構體系班-左程雲

相關文章