如何修改Kadane算法扩展代码以返回k最大子数组的区间？

问题背景 #

给定长度为n的整数数组A，k最大子数组和问题要求找出最多k个不重叠的连续子数组，使它们的和最大；若数组全为负数，结果为0。例如数组[-1,2,-1,2,-1,2,2]的最优两个子数组为[2,-1,2]和[2,2]，总和为7。

已实现的代码能计算最大和，但无法返回构成该和的子数组区间：

• 测试solve_SO([-1,2,-1,2,-1],2)时，期望输出子数组区间[(1,2),(3,4)]
• 测试solve_SO([-1,2,-1,2,-1],1)时，期望输出子数组区间[(1,4)]

原代码如下：

import numpy as np

def solve_SO(test_seq, k=2):
    """
    Computes the k max subarray sum
    """
    num_intervals = k * 2 + 1
    best = np.zeros(num_intervals, dtype=int)
    for seq_idx, val in enumerate(test_seq):
        # Add the current value to all the "include" interval best scores
        for interval_idx in range(1, num_intervals, 2):
            best[interval_idx] += val
        # Go over all intervals from the first include one.
        # If we were better off without, update the best score to be the best
        # from the previous interval. This makes best monotonic.
        # The final value is the current best score overall
        for interval_idx in range(1, num_intervals):
            if best[interval_idx] < best[interval_idx - 1]:
                best[interval_idx] = best[interval_idx - 1]
    return best[num_intervals - 1]

测试示例：

solve_SO([-1, 2, -1, 2, -1], 2)  # 输出4
solve_SO([-1, 2, -1, 2, -1], 1)  # 输出3

修改方案 #

要返回子数组区间，需要在维护最大和的同时，同步跟踪每个状态对应的区间选择。核心是把best数组从单纯的数值，改成存储状态信息字典，包含当前最大和、已选子数组列表、当前正在延续的子数组起止索引（若处于选中子数组的状态）。

修改后的代码 #

def solve_SO_with_intervals(test_seq, k=2):
    n = len(test_seq)
    num_states = 2 * k + 1
    # 初始化状态：每个状态是字典，包含sum(当前最大和)、intervals(已选区间列表)、current_seg(当前延续的子数组起止，None表示未选中)
    best = []
    for i in range(num_states):
        best.append({
            'sum': 0,
            'intervals': [],
            'current_seg': None
        })
    
    for seq_idx, val in enumerate(test_seq):
        # 先复制当前状态，避免更新时覆盖原数据
        prev_best = [s.copy() for s in best]
        
        # 处理所有"选中子数组"的状态（奇数索引：1,3,...,2k-1）
        for state_idx in range(1, num_states, 2):
            prev_state = prev_best[state_idx]
            curr_state = best[state_idx]
            
            if prev_state['current_seg'] is not None:
                # 延续当前子数组：更新和，扩展当前子数组的结束索引
                curr_state['sum'] = prev_state['sum'] + val
                curr_state['current_seg'] = (prev_state['current_seg'][0], seq_idx)
                curr_state['intervals'] = prev_state['intervals'].copy()
            else:
                # 开始新的子数组：比较两种情况：从当前元素开始，或者继承前一个状态的和+当前元素
                option1 = val  # 单独选当前元素
                option2 = prev_best[state_idx - 1]['sum'] + val  # 在前一个状态基础上开始选当前元素
                if option1 > option2:
                    curr_state['sum'] = option1
                    curr_state['current_seg'] = (seq_idx, seq_idx)
                    curr_state['intervals'] = prev_best[state_idx - 1]['intervals'].copy()
                else:
                    curr_state['sum'] = option2
                    curr_state['current_seg'] = (seq_idx, seq_idx)
                    curr_state['intervals'] = prev_best[state_idx - 1]['intervals'].copy()
        
        # 处理所有状态，确保单调性：如果当前状态不如前一个，就继承前一个的状态
        for state_idx in range(1, num_states):
            if best[state_idx]['sum'] < best[state_idx - 1]['sum']:
                # 完全继承前一个状态的所有信息
                best[state_idx]['sum'] = best[state_idx - 1]['sum']
                best[state_idx]['intervals'] = best[state_idx - 1]['intervals'].copy()
                best[state_idx]['current_seg'] = best[state_idx - 1]['current_seg']
        
        # 处理偶数索引的状态（未选中子数组的状态：2,4,...,2k）：这些状态是从选中状态转换而来，需要结束当前子数组
        for state_idx in range(2, num_states, 2):
            prev_selected_state = prev_best[state_idx - 1]
            if prev_selected_state['current_seg'] is not None:
                # 结束当前子数组，加入已选区间列表
                new_intervals = prev_selected_state['intervals'].copy()
                new_intervals.append(prev_selected_state['current_seg'])
                # 比较继承前一个未选中状态，还是结束当前子数组的情况
                if prev_selected_state['sum'] > best[state_idx - 1]['sum']:
                    best[state_idx]['sum'] = prev_selected_state['sum']
                    best[state_idx]['intervals'] = new_intervals
                    best[state_idx]['current_seg'] = None
                else:
                    best[state_idx]['sum'] = best[state_idx - 1]['sum']
                    best[state_idx]['intervals'] = best[state_idx - 1]['intervals'].copy()
                    best[state_idx]['current_seg'] = None
    
    # 处理全负数的情况：如果最大和为0，返回空列表
    final_state = best[-1]
    if final_state['sum'] <= 0:
        return 0, []
    
    # 最后检查是否有未结束的当前子数组，加入区间列表
    if final_state['current_seg'] is not None:
        final_intervals = final_state['intervals'].copy()
        final_intervals.append(final_state['current_seg'])
    else:
        final_intervals = final_state['intervals']
    
    return final_state['sum'], final_intervals

测试验证 #

# 测试用例1：k=2
sum_val, intervals = solve_SO_with_intervals([-1,2,-1,2,-1], 2)
print(sum_val)  # 输出4
print(intervals)  # 输出[(1,2), (3,4)]

# 测试用例2：k=1
sum_val, intervals = solve_SO_with_intervals([-1,2,-1,2,-1], 1)
print(sum_val)  # 输出3
print(intervals)  # 输出[(1,4)]

# 测试全负数情况
sum_val, intervals = solve_SO_with_intervals([-5,-3,-2], 2)
print(sum_val)  # 输出0
print(intervals)  # 输出[]

关键说明 #

1. 状态划分：2k+1个状态对应"未选-选-未选-...-选-未选"的循环，奇数索引表示当前处于选中子数组的状态，偶数索引表示处于未选中的状态。
2. 状态跟踪：每个状态存储已选区间列表和当前正在延续的子数组，确保在更新最大和时同步维护区间信息。
3. 单调性维护：当当前状态的和小于前一个状态时，直接继承前一个状态的所有信息，保证每个状态的和是非递减的。

内容的提问来源于stack exchange，提问作者Simd