【LeetCode 560】和为k的子数组 && 数组三连问题

题目一

题目

给定一个数组和值k，求数组中有多少个子数组，满足和等于值k。

思路

设置ans=0,用于记录答案；
建立字典mp，用于记录前几项累加和信息。具体地，其key为前i+1项的累加和，值为该累加和结果出现过的次数；
从左向右遍历数组nums，不断使用temp_sum计算前i+1项的累加和，并搜索字典是否存在j使得索引为0,...,j的子数组累加和为temp_sum-k，详细见下面的示意图,如果存在，那么我们找到了一个子数组满足题意，把这个累加和为temp_sum-k的子数组的出现次数累加到ans中。
最后，在遍历循环中记录更新累加和字典mp。

代码

def subarraySum(self, nums: List[int], k: int) -> int:
        ans=0
        temp_sum=0
        mp={0:1}
        
        for i in range(len(nums)):
            temp_sum+=nums[i]
            if temp_sum-k in mp.keys():
                ans+=mp[temp_sum-k]
            if temp_sum in mp.keys():
                mp[temp_sum]+=1
            else:
                mp[temp_sum]=1
        return ans

题目2

leetcode plus会员题，链接如下：
https://leetcode.cn/problems/maximum-size-subarray-sum-equals-k/description/
不过，lintcode上不是会员题，所以你可以用下面的链接测试你的代码：
https://www.lintcode.com/problem/911/description

示例：

Input:  nums = [1, -1, 5, -2, 3], k = 3
Output: 4
Explanation:
because the subarray [1, -1, 5, -2] sums to 3 and is the longest.

Input: nums = [-2, -1, 2, 1], k = 1
Output: 2
Explanation:
because the subarray [-1, 2] sums to 1 and is the longest.

代码

说明：这道题可以参考题目一的解法，不同的是，这里定义字典mp的值为累加和对应的下标。

from typing import (
    List,
)

class Solution:
    """
    @param nums: an array
    @param k: a target value
    @return: the maximum length of a subarray that sums to k
    """
    def max_sub_array_len(self, nums: List[int], k: int) -> int:
        # Write your code here
        mp={0:-1} # 注意，这里要自行添加累加和为0时的下标-1，否则答案会跳过那些从索引0开头的子答案。<preSum,idx>
        ans=0
        temp_sum=0
        for i in range(len(nums)):
            temp_sum+=nums[i]
            if temp_sum-k in mp.keys():
                ans=max(ans,i-mp[temp_sum-k])# 更新最长的可行子数组长度
            if temp_sum not in mp.keys():
                mp[temp_sum]=i
        return ans

扩展题目：在二叉树中找到累加和为指定值的最长路径长度

题目链接：
https://www.nowcoder.com/practice/2d35bc3364e3470381bc4eebd9178747

下面的代码，通过了90%的测试算例，不知道怎么进一步优化了。。
这里代码的解法参考了上述数组的解法，只不过这个不是通过索引计算路径长度，而是节点所在层level；另外，在左子树、右子树上都要判断是否存在更好的解。

from operator import le
from sqlite3 import Cursor
import sys

class TreeNode:
    def __init__(self,val=0) -> None:
        self.val=val
        self.left=None
        self.right=None
input_data={}
node_set={}
target=0
i=0
#  create_tree
for line in sys.stdin:    
    a = line.split()
    if i==0:
        num_nodes=eval(a[0])
        head_idx=eval(a[1])
        node_set[head_idx]=TreeNode()
    else:
        if i>num_nodes:
            target=eval(a[0])
            break
        node=node_set[eval(a[0])]
        node.val=eval(a[-1])
        # left child
        left_idx=eval(a[1])
        if left_idx!=0:
            left=TreeNode()
            node_set[left_idx]=left
            node.left=left
        # right child
        right_idx=eval(a[2])
        if right_idx!=0:
            right=TreeNode()
            node_set[right_idx]=right
            node.right=right
    i+=1


head=node_set[head_idx]

def get_max_length(head,target):
    mp={}
    mp[0]=0 # 和为0的所在层为0 # <preSum,level>
    return pre_order(head,target,0,1,0,mp)
def pre_order(head,target,pre_sum,level,max_len,mp):
    if not head:
        return max_len
    cur_sum=pre_sum+head.val
    if cur_sum not in mp.keys():
        mp[cur_sum]=level # 把当前的累加和结果记录到 字典mp中
    if cur_sum-target in mp.keys(): # 如存在，则说明找到一个结果
        max_len=max(max_len,level-mp[cur_sum-target])
    
    max_len=pre_order(head.left,target,cur_sum,level+1,max_len,mp)
    max_len=pre_order(head.right,target,cur_sum,level+1,max_len,mp)

    if level==mp[cur_sum]: # 深度优先遍历，清除访问痕迹
        del mp[cur_sum]
    return max_len
if __name__=='__main__':
    print(get_max_length(head,target))