前K个高频元素3种解题方法

题目

给定一个非空的整数数组，返回其中出现频率前 k 高的元素。
示例 1:

输入: nums = [1,1,1,2,2,3], k = 2
输出: [1,2]

示例 2:

输入: nums = [1], k = 1
输出: [1]

提示：

你可以假设给定的 k 总是合理的，且 1 ≤ k ≤ 数组中不相同的元素的个数。
你的算法的时间复杂度必须优于 O(n log n) , n 是数组的大小。
题目数据保证答案唯一，换句话说，数组中前 k 个高频元素的集合是唯一的。
你可以按任意顺序返回答案。

方法一 优先队列

Swift代码

final class FixedSizePriorityQueue<Element> {
    
    let capacity: Int
    
    fileprivate(set) var count: Int
    
    var isEmpty: Bool { count == 0 }
    
    private var pq, qp: [Int]
    
    var elements: [Element?]
    
    let comparator: ((Element, Element)-> Bool)
    
    public init(_ capacity: Int, comparator: @escaping ((Element, Element)-> Bool)) {
        self.capacity = capacity
        count = 0
        pq = Array(repeating: 0, count: capacity + 1)
        qp = Array(repeating: -1, count: capacity + 1)
        elements = Array(repeating: nil, count: capacity + 1)
        self.comparator = comparator
    }
    
    public func contains(_ i: Int) -> Bool {
        return qp[i] != -1
    }
    
    public func insert(_ i: Int, _ element: Element) {
        count += 1
        qp[i] = count
        pq[count] = i
        elements[i] = element
        swim(count)
    }
    
    public func topIndex() -> Int {
        return pq[1]
    }
    
    public func topElement() -> Element? {
        return elementOf(topIndex())
    }
    
    public func deleteTop() -> Element? {
        if isEmpty {
            return nil
        }
        let element = topElement()
        let min = pq[1]
        exchange(1, count)
        count -= 1
        sink(1)
        qp[min] = -1
        elements[min] = nil
        pq[count + 1] = -1
        return element
    }
    
    public func elementOf(_ i: Int) -> Element? {
        return elements[i]
    }
    
    public func changeKey(_ i: Int, _ key: Element) {
        elements[i] = key
        swim(qp[i])
        sink(qp[i])
    }
    
    public func compare(_ i: Int, _ j: Int) -> Bool {
        return comparator(elements[pq[i]]!, elements[pq[j]]!)
    }
    
    public func delete(_ i: Int) {
        let index = qp[i]
        exchange(index, count)
        count -= 1
        swim(index)
        sink(index)
        elements[i] = nil
        qp[i] = -1
    }
    
    public func exchange(_ i: Int, _ j: Int) {
        pq.swapAt(i, j)
        qp[pq[i]] = i
        qp[pq[j]] = j
    }
    
    private func swim(_ k: Int) {
        var k = k
        while k > 1 && compare(k >> 1, k) {
            pq.swapAt(k >> 1, k)
            k >>= 1
        }
    }

    private func sink(_ k: Int) {
        var k = k, j = k << 1
        while j <= count {
            if j < count && compare(j, j + 1) { j += 1 }
            if !compare(k, j) { break }
            pq.swapAt(k, j)
            k = j
            j = k << 1
        }
    }
    
}

class Solution {
    func topKFrequent(_ nums: [Int], _ k: Int) -> [Int] {
    var dic = [Int: Int]()
    for item in nums {
        if let value = dic[item] {
            dic[item] = value + 1
        } else {
            dic[item] = 1
        }
    }
    let queue = FixedSizePriorityQueue<(Int, Int)>(dic.count, comparator: { $0.1 < $1.1 })
    for item in dic.enumerated() {
        queue.insert(item.offset, item.element)
    }
    var result = [Int]()
    while result.count < k && queue.count > 0 {
        result.append(queue.deleteTop()!.0)
    }
    
    return result
}
}

方法二 红黑树

Swift代码

class RedBlackTree<Key: Comparable, Value> {
    
    final class Node {
        
        public var item: (Key, Value)
        
        public var left: Node?
        
        public var right: Node?
        
        public var count: Int
        
        public var color: Color
        
        public init(
            key: Key,
            value: Value,
            left: Node? = nil,
            right: Node? = nil,
            count: Int = 1,
            color: Color = .red
        ) {
            self.item = (key, value)
            self.left = left
            self.right = right
            self.count = count
            self.color = color
        }
        
        public func rotateLeft() -> Node? {
            guard let x = right else {
                return nil
            }
            right = x.left
            x.left = self
            x.color = x.left?.color ?? .black
            x.left?.color = .red
            x.count = count
            count = (left?.count ?? 0) + (right?.count ?? 0) + 1
            return x
        }
        
        public func rotateRight() -> Node? {
            guard let x = left else {
                return nil
            }
            left = x.right
            x.right = self
            x.color = x.right?.color ?? .black
            x.right?.color = .red
            x.count = count
            count = (left?.count ?? 0) + (right?.count ?? 0) + 1
            return x
        }
        
        public func flipColors() {
            color.toggle()
            left?.color.toggle()
            right?.color.toggle()
        }
        private var stop = false
        func reverseMiddleOrder(_ block: @escaping ((Key, Value), inout Bool) -> Void) {
            stop = false
            right?.reverseMiddleOrder(block)
            block(item, &stop)
            if !stop {
                left?.reverseMiddleOrder(block)
            }
        }
    }
    
    var root: Node? = nil
    
    var count: Int { size(root) }
    
    subscript(key: Key) -> Value? {
        set {
            put(key, newValue)
        }
        get {
            get(key)
        }
    }
    
    private func isRed(_ node: Node?) -> Bool {
        node?.color == .red
    }
    
    private func size(_ node: Node?) -> Int {
        node?.count ?? 0
    }
    
    public var isEmpty: Bool {
        root == nil
    }
    
    public func get(_ key: Key) -> Value? {
        return get(root, key)
    }
    
    private func get(_ node: Node?, _ key: Key) -> Value? {
        var node = node
        while node != nil {
            if key < node!.item.0 {
                node = node?.left
            } else if key > node!.item.0 {
                node = node?.right
            } else {
                return node?.item.1
            }
        }
        return nil
    }
    
    public func contains(_ key: Key) -> Bool {
        get(key) != nil
    }
    
    public func put(_ key: Key, _ value: Value?) {
        if value == nil {
            return
        }
        root = put(root, key, value!)
        root?.color = .black
    }
    
    private func put(_ node: Node?, _ key: Key, _ value: Value) -> Node? {
        guard node != nil else {
            return Node(key: key, value: value)
        }
        var node = node
        if key < node!.item.0 {
            node?.left = put(node?.left, key, value)
        } else if key > node!.item.0 {
            node?.right = put(node?.right, key, value)
        } else {
            node?.item.1 = value
        }
        
        if isRed(node?.right) && !isRed(node?.left) {
            node = node?.rotateLeft()
        }
        if isRed(node?.left) && isRed(node?.left?.left) {
            node = node?.rotateRight()
        }
        if isRed(node?.left) && isRed(node?.right) {
            node?.flipColors()
        }
        
        node?.count = size(node?.left) + size(node?.right) + 1
        return node
    }
    
    func balance(_ node: Node?) -> Node? {
        var node = node
        if isRed(node?.right) {
            node = node?.rotateLeft()
        }
        if isRed(node?.left) && isRed(node?.left?.left) {
            node = node?.rotateRight()
        }
        if isRed(node?.left) && isRed(node?.right) {
            node?.flipColors()
        }
        node?.count = size(node?.left) + size(node?.right) + 1
        return node
    }
    
    enum Color {
        
        case red
        
        case black
        
        mutating func toggle() {
            switch self {
            case .red:
                self = .black
            case .black:
                self = .red
            }
        }
    }
}

class Solution {
    func topKFrequent(_ nums: [Int], _ k: Int) -> [Int] {
        var counter = Dictionary<Int, Int>()
        for item in nums {
            if let value = counter[item] {
                counter[item] = value + 1
            } else {
                counter[item] = 1
            }
        }
        let tree = RedBlackTree<Int, [Int]>()
        for item in counter {
            var items = tree[item.1]
            if items == nil {
                items = [item.0]
            } else {
                items?.append(item.0)
            }
            tree[item.1] = items
        }
        var items = [Int]()
        tree.root?.reverseMiddleOrder { (item, stop)  in
            if items.count < k {
                for element in item.1 {
                    items.append(element)
                    if items.count == k { break }
                }
            } else {
                stop = true
            }
        }
        return items
    }
}

方法三 排序

排序并选取前k元素。该版本的快速排序比系统排序快一些。

extension Int {
    mutating func increasing() -> Int {
        self += 1
        return self
    }
    mutating func decreasing() -> Int {
        self -= 1
        return self
    }
}
extension Array {
    mutating func quickSort(by comparator: @escaping (Element, Element) -> Bool ) {
        quickSort(low: 0, high: count - 1, by: comparator)
    }
    mutating func quickSort(low: Int, high: Int, by comparator: @escaping (Element, Element) -> Bool ) {
        if high <= low { return }
        let j = partition(low: low, high:high, by: comparator)
        quickSort(low: low, high: j - 1, by: comparator)
        quickSort(low: j + 1, high: high, by: comparator)
    }

    mutating func partition(low: Int, high: Int, by comparator: @escaping (Element, Element) -> Bool) -> Int {
        var i = low
        var j = high + 1
        let v = self[low]
        while true {
            while comparator(self[i.increasing()], v) && i != high {}
            while comparator(v, self[j.decreasing()]) && i != low {}
            if i >= j { break }
            swapAt(i, j)
        }
        swapAt(low, j)
        return j
    }
}

class Solution {
    func topKFrequent(_ nums: [Int], _ k: Int) -> [Int] {
        var items = nums.reduce(into: [Int: Int]()) { (counter, item) in
            counter[item] = (counter[item] ?? 0) + 1
        }.map { $0 }
        items.quickSort {
            $0.1 > $1.1
        }
        return items[0..<k].map { $0.0 }
    }
}