数据结构与算法之堆: LeetCode 215. 数组中的第K个最大元素 (Ts, Py, Go版)

本文链接：https://blog.csdn.net/Tyro_java/article/details/133471440

数组中的第K个最大元素

https://leetcode.cn/problems/kth-largest-element-in-an-array/

描述

给定整数数组 nums 和整数 k，请返回数组中第 k 个最大的元素。
请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。
你必须设计并实现时间复杂度为 O(n) 的算法解决此问题。

示例 1

输入: [3,2,1,5,6,4], k = 2
输出: 5

示例 2

输入: [3,2,3,1,2,4,5,5,6], k = 4
输出: 4

提示

1 <= k <= nums.length <= $10^5$
- $10^4$ <= nums[i] <= $10^4$

Typescript 版算法实现

1 ）基于js中原生sort api

const findKthLargest = function(nums, k) {
    return nums.sort((a,b) => b-a)[k - 1]
};

这个浏览器默认提供的sort()方法，一般时间复杂度是 O(nlogn)

2 ）基于堆的数据结构和堆排序的方法

// 建立最小堆类
class MinHeap {
    heap: number[] = [];
    // 交换节点位置
    swap(i, j) {
        [this.heap[i], this.heap[j]] = [this.heap[j], this.heap[i]];
    }
    // 获得父节点
    getParentIndex(i) {
        return (i - 1) >> 1;
    }
    // 获取左子节点
    getLeftIndex(i) {
        return (i << 1) + 1; // 极客写法
    }
    // 获取右子节点
    getRightIndex(i) {
        return (i << 1) + 2;
    }
    // 向上移动
    shiftUp(index) {
        // 如果到了堆顶元素，index是0，则不要再上移了
        if(!index) return;
        const parentIndex = this.getParentIndex(index)
        if(this.heap[parentIndex] > this.heap[index]) {
            this.swap(parentIndex, index)
            this.shiftUp(parentIndex)
        }
    }
    // 下移
    shiftDown(index) {
        // 边界1：如果到了堆尾元素，则不要再下移了
        if(index >= this.heap.length - 1) return;
        const size = this.size();
        const leftIndex = this.getLeftIndex(index);
        const rightIndex = this.getRightIndex(index);
        if (leftIndex < size && this.heap[leftIndex] < this.heap[index]) {
            this.swap(leftIndex, index);
            this.shiftDown(leftIndex);
        }
        if (rightIndex < size && this.heap[rightIndex] < this.heap[index]) {
            this.swap(rightIndex, index);
            this.shiftDown(rightIndex);
        }
    }
    // 插入
    insert(value) {
        this.heap.push(value);
        this.shiftUp(this.heap.length - 1);
    }
    // 删除堆顶
    pop() {
        // pop()方法删除数组最后一个元素并返回，赋值给堆顶
        this.heap[0] = this.heap.pop();
        // 对堆顶重新排序
        this.shiftDown(0);
    }
    // 获取堆顶
    peak() {
        return this.heap[0];
    }
    // 获取堆的大小
    size() {
        return this.heap.length;
    }
}

// 实现
const findKthLargest = (nums, k) => {
    const h = new MinHeap();
    nums.forEach(n => {
        // 将数组元素依次插入堆中
        h.insert(n);
        // 如果堆满，则执行优胜劣汰
        (h.size() > k) && h.pop();
    })
    // 返回堆顶，此时就是第k大的元素
    return h.peak();
};

关键在于这个堆的数据结构提供的 insert 方法与 pop 方法
时间复杂度：O(nlogk)
- 一个n循环，里面还嵌套一个heap的上移递归操作logk
- 总体：n*logk
空间复杂度: O(k) 或 O(logn)
- 堆的大小，数组的大小, k是输入的堆大小
注意
- 本题使用的是一个堆排序的算法，O(nlogn)
- 但是还有其他排序也可以达到这个效率
- 但是这个不符合题目的要求：时间复杂度为 O(n) 的算法

3 ）基于冒泡排序

function findKthLargest(nums: number[], k: number): number {
    const len = nums.length - 1;
    for (let i = len; i > len - k; i--) {
        for (let j = 0; j < i; j++) {
            if (nums[j] > nums[j + 1]) {
                [nums[j], nums[j + 1]] = [nums[j + 1], nums[j]];
            }
        }
    }
    return nums[len - k + 1];
};

这个算法基于冒泡排序修改而来，但是无法通过 leetcode , 对于一些用例数据，超出时间限制

4 ）基于快速排序

function quickselect(nums: number[], l: number, r: number, k: number) {
        if (l == r) return nums[k];
        const pivot = nums[l];
        let i:number = l - 1, j:number = r + 1;
        while (i < j) {
            do i++; while (nums[i] < pivot);
            do j--; while (nums[j] > pivot);
            if (i < j) [nums[i], nums[j]] = [nums[j], nums[i]];
        }
        return k <= j ? quickselect(nums, l, j, k) : quickselect(nums, j + 1, r, k);
    }

function findKthLargest(nums: number[], k: number): number {
    const n = nums.length;
    return quickselect(nums, 0, n - 1, n - k);
}

这是官方提供的方法
时间复杂度 O(n), 证明过程可以参考「《算法导论》9.2：期望为线性的选择算法」
空间复杂度 O(logn)，递归使用栈空间的空间代价的期望为 O(log⁡n)O(\log n)O(logn)

Python3 版算法实现

1 ) 方案1

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        nums.sort(reverse=True)
        return nums[k - 1]

2 ）方案2

class MinHeap:
    def __init__(self):
        self.heap = []

    def swap(self, i, j):
        self.heap[i], self.heap[j] = self.heap[j], self.heap[i]

    def get_parent_index(self, i):
        return (i - 1) // 2

    def get_left_index(self, i):
        return (i << 1) + 1

    def get_right_index(self, i):
        return (i << 1) + 2

    def shift_up(self, index):
        if index == 0:
            return
        parent_index = self.get_parent_index(index)
        if self.heap[parent_index] > self.heap[index]:
            self.swap(parent_index, index)
            self.shift_up(parent_index)

    def shift_down(self, index):
        if index >= len(self.heap) - 1:
            return
        size = len(self.heap)
        left_index = self.get_left_index(index)
        right_index = self.get_right_index(index)
        if left_index < size and self.heap[left_index] < self.heap[index]:
            self.swap(left_index, index)
            self.shift_down(left_index)
        if right_index < size and self.heap[right_index] < self.heap[index]:
            self.swap(right_index, index)
            self.shift_down(right_index)

    def insert(self, value):
        self.heap.append(value)
        self.shift_up(len(self.heap) - 1)

    def pop(self):
        if len(self.heap) <= 1:
            return
        self.heap[0] = self.heap.pop()
        self.shift_down(0)

    def peak(self):
        return self.heap[0]

    def size(self):
        return len(self.heap)

class Solution:
    def findKthLargest(self, nums, k):
        h = MinHeap()
        for n in nums:
            h.insert(n)
            if h.size() > k:
                h.pop()
        return h.peak()

Leetcode 超出时间限制

3 ) 方案3

class Solution:
    def findKthLargest(self, nums, k):
        length = len(nums) - 1
        for i in range(length, length - k, -1):
            for j in range(i):
                if nums[j] > nums[j + 1]:
                    nums[j], nums[j + 1] = nums[j + 1], nums[j]
        return nums[length - k + 1]

Leetcode 超出时间限制

4 ）方案4

class Solution:
    def quickselect(self, nums, l, r, k):
        if l == r:
            return nums[k]
        pivot = nums[l]
        i, j = l - 1, r + 1
        while i < j:
            i += 1
            while nums[i] < pivot:
                i += 1
            j -= 1
            while nums[j] > pivot:
                j -= 1
            if i < j:
                nums[i], nums[j] = nums[j], nums[i]
        if k <= j:
            return self.quickselect(nums, l, j, k)
        return self.quickselect(nums, j + 1, r, k)

    def findKthLargest(self, nums: List[int], k: int) -> int:
        n = len(nums)
        return self.quickselect(nums, 0, n - 1, n - k)

Golang 版算法实现

1 ) 方案1

func findKthLargest(nums []int, k int) int {
	sort.Sort(sort.Reverse(sort.IntSlice(nums)))
	return nums[k-1]
}

2 ）方案2

type MinHeap struct {
	heap []int
}

// 交换节点位置
func (h *MinHeap) swap(i, j int) {
	h.heap[i], h.heap[j] = h.heap[j], h.heap[i]
}

// 获得父节点
func (h *MinHeap) getParentIndex(i int) int {
	return (i - 1) >> 1
}

// 获取左子节点
func (h *MinHeap) getLeftIndex(i int) int {
	return (i << 1) + 1
}

// 获取右子节点
func (h *MinHeap) getRightIndex(i int) int {
	return (i << 1) + 2
}

// 向上移动
func (h *MinHeap) shiftUp(index int) {
	if index == 0 {
		return
	}
	parentIndex := h.getParentIndex(index)
	if h.heap[parentIndex] > h.heap[index] {
		h.swap(parentIndex, index)
		h.shiftUp(parentIndex)
	}
}

// 下移
func (h *MinHeap) shiftDown(index int) {
	if index >= len(h.heap)-1 {
		return
	}
	size := h.size()
	leftIndex := h.getLeftIndex(index)
	rightIndex := h.getRightIndex(index)
	if leftIndex < size && h.heap[leftIndex] < h.heap[index] {
		h.swap(leftIndex, index)
		h.shiftDown(leftIndex)
	}
	if rightIndex < size && h.heap[rightIndex] < h.heap[index] {
		h.swap(rightIndex, index)
		h.shiftDown(rightIndex)
	}
}

// 插入
func (h *MinHeap) insert(value int) {
	h.heap = append(h.heap, value)
	h.shiftUp(len(h.heap) - 1)
}

// 删除堆顶
func (h *MinHeap) pop() {
	if len(h.heap) <= 1 {
		return
	}
	h.heap[0] = h.heap[len(h.heap)-1]
	h.heap = h.heap[:len(h.heap)-1]
	h.shiftDown(0)
}

// 获取堆顶
func (h *MinHeap) peak() int {
	return h.heap[0]
}

// 获取堆的大小
func (h *MinHeap) size() int {
	return len(h.heap)
}

// 实现
func findKthLargest(nums []int, k int) int {
	h := &MinHeap{}
	for _, n := range nums {
		h.insert(n)
		if h.size() > k {
			h.pop()
		}
	}
	return h.peak()
}

在 Leetcode 上超出时间限制

3 ) 方案3

func findKthLargest(nums []int, k int) int {
	len := len(nums) - 1
	for i := len; i > len-k; i-- {
		for j := 0; j < i; j++ {
			if nums[j] > nums[j+1] {
				nums[j], nums[j+1] = nums[j+1], nums[j]
			}
		}
	}
	return nums[len-k+1]
}

在 Leetcode 上超出时间限制

4 ）方案4

func quickselect(nums []int, l, r, k int) int {
	if l == r {
		return nums[k]
	}
	pivot := nums[l]
	i, j := l-1, r+1
	for i < j {
		i++
		for nums[i] < pivot {
			i++
		}
		j--
		for nums[j] > pivot {
			j--
		}
		if i < j {
			nums[i], nums[j] = nums[j], nums[i]
		}
	}
	if k <= j {
		return quickselect(nums, l, j, k)
	}
	return quickselect(nums, j+1, r, k)
}

func findKthLargest(nums []int, k int) int {
	n := len(nums)
	return quickselect(nums, 0, n-1, n-k)
}