面试算法题：判断数组前n-2个数加减组合是否等于a mod b

First, let's restate the problem clearly to make sure we're aligned:

Given an integer array where the last two elements are a and b, determine if we can assign a + or - sign to every element except the last two such that their combined sum equals a mod b.

Step 1: Handle Critical Edge Cases #

Before diving into core logic, we need to address a few edge cases upfront:

• Invalid Input: If the array has fewer than 2 elements, we can't extract a and b—return false (or throw an error, depending on your requirements).
• Undefined Modulo: If b is 0, a mod b is mathematically undefined—return false.
• No Elements to Combine: If the array has exactly 2 elements (only a and b), we need to check if 0 equals a mod b (since there are no elements to sum).

Step 2: Compute the Target Value #

First, calculate target = a mod b. Note that many programming languages return negative results for modulo operations with negative a. To get the non-negative remainder (aligned with standard mathematical definitions), adjust it using:

target = ((a % b) + b) % b

This ensures target falls in the range [0, |b| - 1] regardless of the signs of a and b.

Step 3: Check if ± Combinations Can Form the Target #

This is a classic dynamic programming problem where we track all possible sums achievable by adding or subtracting each element in the list.

Approach Explanation

• Start with a set of possible sums initialized to {0} (since the sum of zero elements is 0).
• For each element in the list (excluding a and b), update the set by adding both the element and its negative to every existing sum.
• After processing all elements, check if the target exists in the final set of possible sums.

Code Example (Python)

def can_form_target(arr, target):
    possible_sums = {0}
    for num in arr:
        new_sums = set()
        for s in possible_sums:
            new_sums.add(s + num)
            new_sums.add(s - num)
        possible_sums = new_sums
    return target in possible_sums

def solve(nums):
    # Edge case: Not enough elements to extract a and b
    if len(nums) < 2:
        return False
    
    a, b = nums[-2], nums[-1]
    
    # Edge case: Modulo by zero is undefined
    if b == 0:
        return False
    
    # Compute non-negative target value
    target = ((a % b) + b) % b
    
    # Get elements to combine (all except last two)
    elements = nums[:-2]
    
    # Edge case: No elements to combine—check if target is 0
    if not elements:
        return target == 0
    
    return can_form_target(elements, target)

Optimizations for Larger Inputs #

The set-based approach works well for small to medium arrays, but for larger inputs (30+ elements), the number of possible sums grows exponentially (up to 2^n). Here’s how to optimize:

• Boolean Array: Calculate the minimum and maximum possible sums ( -sum(abs(x) for x in elements) and sum(abs(x) for x in elements) ). Use a boolean array where each index represents a sum (shifted by the minimum sum to avoid negative indices) for faster lookups and better memory efficiency.
• Early Exit: If the target is found in the possible sums at any step, you can immediately return true without processing remaining elements.

Example Walkthrough #

Let’s test with nums = [3, -2, 1, 5]:

• a = 1, b =5 → target = 1 mod5 =1.
• Elements to combine: [3, -2].
• Possible sums after processing 3: {3, -3}.
• Possible sums after processing -2: {1,5,-5,-1}.
• Target 1 is present → return true.

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