来自Leetcode的问题如下:给定一个整数数组,计算该数组的枢轴索引
枢轴索引是严格在索引左边的所有数的和等于严格在索引右边的所有数的和的索引。
如果索引位于数组的左边缘,则左和为0,因为左边没有元素。这也适用于数组的右边缘。
返回最左边的枢轴索引。如果不存在这样的索引,返回-1。
示例1:
输入:nums = [1,7,3,6,5,6]
输出:3
解释:主指数是3。左sum = nums[0] + nums[1] + nums[2] = 1 + 7 + 3 = 11右sum = nums[4] + nums[5] = 5 + 6 = 11
示例2:
输入:nums = [1,2,3]
输出:1
解释:没有索引满足问题语句中的条件。
示例3:
输入:nums = [2,1,-1]
输出:0
解释:主指数是0。左sum = 0(索引0的左边没有元素)右sum = nums[1] + nums[2] = 1 + -1 = 0
我的解决方案有效,但超过了Leetcode的时间限制。达到测试用例740/745我该如何进一步优化代码?
class Solution:
def pivotIndex(self, nums: List[int]) -> int:
left_list = []
for index,num in enumerate(nums):
if sum(nums[index+1:]) == sum(left_list):
return index
else:
left_list.append(num)
return -1
不断地对列表中的切片求和将是非常低效的。你需要一个不同的策略。像这样:
class Solution:
def pivotIndex(self, nums: List[int]) -> int:
total = sum(nums)
left = 0
for i, x in enumerate(nums):
if left == total - left - x:
return i
left += x
return -1
你不需要做多次求和,你只需要做一个累积求和,然后每次迭代你可以用它来计算右半部分的和减去左边的和。
from itertools import accumulate
from operator import add
from typing import List
class Solution:
def pivotIndex(self, nums: List[int]) -> int:
cumsum = accumulate(nums, add)
final_sum = sum(nums)
last_sum = 0 # to skip summing the current element in evaluation.
for index, left_sum in enumerate(cumsum):
right_sum = final_sum - left_sum
if last_sum == right_sum:
return index
last_sum = left_sum
return -1
inputs = [1,7,3,6,5,6]
print(Solution().pivotIndex(inputs))
# 3