两个有序数组中第K大的数

一道算法上机题目，大概记录下思路。

很容易想到的一个思路是，将两个数组归并，但这需要常数级的时间复杂度。

两个数组A和B，各自有序，寻找其中第K大的数。

可以从数组A中选前i个数，从数组B中选前j个数，使 i + j = k

再去比较一下 A[i] 和 B[j] 的大小关系，如此一来，便可以排除一部分元素。

• A[i] < B[j]：我们的目标便锁定在A[i+1:]和B[:j+1]之间，接着在剩下的区间里找第 k - i 大的数
• A[i] > B[j]：我们的目标便锁定在A[:i+1]和B[j+1:]之间，接着在剩下的区间里找第 k - j 大的数

package main

import (
	"fmt"
	"math/rand"
	"sort"
	"strconv"
)

func main() {
	var nums1, nums2 []int
	set := make(map[int]bool)
	for i := 0; i < 10; i++ {
		randNum := rand.Intn(100)
		if _, ok := set[randNum]; !ok {
			set[randNum] = true
			nums1 = append(nums1, randNum)
		}

		randNum = rand.Intn(100)
		if _, ok := set[randNum]; !ok {
			set[randNum] = true
			nums2 = append(nums2, randNum)
		}
	}
	nums := append(nums1, nums2...)
	sort.Ints(nums1)
	sort.Ints(nums2)
	sort.Ints(nums)
	fmt.Println(nums1, nums2)
	fmt.Println(nums)

	fmt.Println("expect:" + strconv.Itoa(nums[8]))
	fmt.Println("what I get:" + strconv.Itoa(kthLargestNum(nums1, nums2, 9)))
}

func kthLargestNum(A, B []int, k int) int {
	if len(A) > len(B) {
		return kthLargestNum(B, A, k)
	}
	if len(A) == 0 {
		return B[k-1]
	}
	if k == 1 {
		return min(A[0], B[0])
	}

	var i, j int
	// average if possible
	if k/2 < len(A) {
		i = (k / 2) - 1
	} else {
		i = len(A) - 1
	}
	j = k - (i + 1) - 1

	fmt.Println(A, B, k, i, j)

	if A[i] < B[j] {
		return kthLargestNum(A[i+1:], B[:j+1], k-i-1)
	} else {
		return kthLargestNum(A[:i+1], B[j+1:], k-j-1)
	}
}

func min(a, b int) int {
	if a > b {
		return b
	}
	return a
}