如何在NumPy中高效求两个排序数组的集合差？

Hey there! Let's break down your two NumPy questions with efficient, tailored solutions—especially focusing on that O(n+m) requirement for the second one.

Since both arrays are already sorted, we can skip the expensive sorting steps that come with generic set operations. Here are two solid approaches:

• Option 1: Use NumPy's built-in optimized function
  The np.setdiff1d function calculates the set difference (elements in the first array not present in the second), and you can enable optimizations by flagging that inputs are sorted and unique:

  import numpy as np
  
  sorted_arr1 = np.array([1, 3, 5, 7, 9])
  sorted_arr2 = np.array([3, 7])
  
  # Leverage sorted/unique inputs for linear-time processing
  diff = np.setdiff1d(sorted_arr1, sorted_arr2, assume_unique=True, sorted=True)
  print(diff)  # Output: [1 5 9]
  

  The assume_unique=True tells NumPy there are no duplicates in either array, and sorted=True skips pre-sorting, making this run in O(n+m) time.

• Option 2: Manual two-pointer method (great for duplicates)
  If your arrays might have duplicates (but are still sorted), a two-pointer approach gives you full control and maintains linear time complexity:

  def sorted_set_diff(arr1, arr2):
      i = j = 0
      result = []
      n, m = len(arr1), len(arr2)
      
      while i < n and j < m:
          if arr1[i] < arr2[j]:
              result.append(arr1[i])
              i += 1
          elif arr1[i] > arr2[j]:
              j += 1
          else:
              # Skip matching elements
              i += 1
              j += 1
      # Add any remaining elements from arr1
      result.extend(arr1[i:])
      return np.array(result)
  
  arr1 = np.array([1, 3, 3, 5, 7])
  arr2 = np.array([3, 6, 7])
  print(sorted_set_diff(arr1, arr2))  # Output: [1 3 5]
  

Since both rows and exclude are sorted, the two-pointer method is the perfect choice here—it’s strictly O(n+m) time, beats slower approaches like np.isin (O(n log m)) or set conversions (which messes up order and wastes cycles on sorting).

Here’s the implementation:

import numpy as np

def filter_sorted_rows(rows, exclude):
    i = j = 0
    n, m = len(rows), len(exclude)
    filtered = []
    
    while i < n and j < m:
        if rows[i] < exclude[j]:
            filtered.append(rows[i])
            i += 1
        elif rows[i] > exclude[j]:
            j += 1
        else:
            # Skip the index we need to exclude
            i += 1
            j += 1
    # Add any remaining rows that weren't excluded
    filtered.extend(rows[i:])
    return np.array(filtered)

# Test it out
rows = np.array([0, 1, 2, 3, 4, 5, 6])
exclude = np.array([2, 4, 5])
print(filter_sorted_rows(rows, exclude))  # Output: [0 1 3 6]

Why this works so well: #

• No nested loops or sorting—just a single pass through both arrays
• Preserves the original sorted order of rows (critical if you need to keep indexing intact)
• Minimal memory overhead, since we only build the result array as we go

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