题目描述:爱丽丝和鲍勃有不同大小的糖果棒:A[i] 是爱丽丝拥有的第 i 根糖果棒的大小,B[j] 是鲍勃拥有的第 j 根糖果棒的大小。
因为他们是朋友,所以他们想交换一根糖果棒,这样交换后,他们都有相同的糖果总量。(一个人拥有的糖果总量是他们拥有的糖果棒大小的总和。)
返回一个整数数组 ans,其中 ans[0] 是爱丽丝必须交换的糖果棒的大小,ans[1] 是 Bob 必须交换的糖果棒的大小。
如果有多个答案,你可以返回其中任何一个。保证答案存在。
解题思路:因为一定有可行解,那么我们可以知道可行解之差肯定为两个人原来糖果总数之差的二分之一,只有这样在交换之后才能拿到相同的糖果数,我们只需要找出一对这样的糖果大小即可,为了减小时间复杂度,先对两个数组都进行排序,这样在查找的时候就可以同时往一个方向查找,不需要遍历两个数组中的所有数对,代码如下:
class Solution:
def fairCandySwap(self, A: List[int], B: List[int]) -> List[int]:
A.sort()
B.sort()
suma = sum(A)
sumb = sum(B)
lena, lenb = len(A), len(B)
if suma < sumb:
diff = (sumb - suma) // 2
j = 0
for i in range(lena):
while(j < lenb and B[j] < A[i]+diff):
j += 1
if A[i] + diff == B[j]:
return [A[i], B[j]]
else:
diff = (suma - sumb) // 2
j = lenb - 1
for i in range(lena-1, -1, -1):
while(j >= 0 and B[j]+diff > A[i]):
j -= 1
if B[j] + diff == A[i]:
return [A[i], B[j]]