Learnixo

Python Essentials for AI Engineers · Lesson 33 of 36

List and Array Problems

Problem 1: Two Sum

Problem: Given a list of numbers and a target, return the indices of two numbers that add up to the target. Assume exactly one solution exists.

Python
def two_sum(nums: list[int], target: int) -> tuple[int, int]:
    """
    Find two indices whose values sum to target.
    O(n) time using a hash map — much better than O(n²) brute force.
    """
    seen: dict[int, int] = {}   # value  index

    for i, num in enumerate(nums):
        complement = target - num
        if complement in seen:
            return (seen[complement], i)
        seen[num] = i

    raise ValueError("No solution found")


print(two_sum([2, 7, 11, 15], 9))    # (0, 1)  2 + 7 = 9
print(two_sum([3, 2, 4], 6))         # (1, 2)  2 + 4 = 6
print(two_sum([3, 3], 6))            # (0, 1)

Why it works: For each number, we check if the complement (target - num) was already seen. If yes, we have our pair. We only need one pass through the list.


Problem 2: Sliding Window Maximum

Problem: Given an array and window size k, return the maximum value in each sliding window of size k.

Python
from collections import deque

def sliding_window_max(nums: list[float], k: int) -> list[float]:
    """
    Find the max in each window of size k.
    O(n) time using a monotonic deque.

    AI context: useful for finding peak confidence scores in token sequences.
    """
    if not nums or k <= 0:
        return []

    result = []
    dq: deque[int] = deque()   # Stores indices, largest values at front

    for i, val in enumerate(nums):
        # Remove indices that are out of the current window
        while dq and dq[0] < i - k + 1:
            dq.popleft()

        # Remove indices whose values are smaller than current
        # (they can never be the max for any future window)
        while dq and nums[dq[-1]] < val:
            dq.pop()

        dq.append(i)

        # Start recording results once the first full window is complete
        if i >= k - 1:
            result.append(nums[dq[0]])

    return result


# Confidence scores from a classifier over a document
scores = [0.9, 0.4, 0.85, 0.7, 0.92, 0.3, 0.88]
print(sliding_window_max(scores, k=3))
# [0.9, 0.85, 0.92, 0.92, 0.92]

Deque invariant: The deque stores indices in decreasing order of their values. The front is always the index of the maximum.


Problem 3: Merge Two Sorted Arrays

Problem: Given two sorted arrays, merge them into a single sorted array without using sorted().

Python
def merge_sorted(a: list[float], b: list[float]) -> list[float]:
    """
    Merge two sorted arrays into one sorted array.
    O(m + n) time, O(m + n) space.

    AI context: merging sorted retrieval results from two different vector stores.
    """
    result = []
    i = j = 0

    while i < len(a) and j < len(b):
        if a[i] <= b[j]:
            result.append(a[i])
            i += 1
        else:
            result.append(b[j])
            j += 1

    # One of these is a no-op (the exhausted list)
    result.extend(a[i:])
    result.extend(b[j:])

    return result


store_a_scores = [0.45, 0.72, 0.88, 0.95]
store_b_scores = [0.30, 0.65, 0.91, 0.97]
merged = merge_sorted(store_a_scores, store_b_scores)
print(merged)
# [0.3, 0.45, 0.65, 0.72, 0.88, 0.91, 0.95, 0.97]

Problem 4: Find All Duplicates

Problem: Given a list of integers in the range 1 to n (where n is the list length), find all elements that appear twice.

Python
def find_duplicates(nums: list[int]) -> list[int]:
    """
    Find all duplicates in O(n) time, O(1) extra space.
    
    Trick: use the sign of nums[abs(num) - 1] as a visited flag.
    This only works because values are in range [1, n].
    """
    duplicates = []

    for num in nums:
        idx = abs(num) - 1   # Map value to index
        if nums[idx] < 0:
            # Already flipped  this value is a duplicate
            duplicates.append(abs(num))
        else:
            nums[idx] = -nums[idx]   # Mark as visited

    # Restore original (optional)
    for i in range(len(nums)):
        nums[i] = abs(nums[i])

    return duplicates


print(find_duplicates([4, 3, 2, 7, 8, 2, 3, 1]))   # [2, 3]
print(find_duplicates([1, 1, 2]))                    # [1]

Alternative for general inputs: Use a Counter or set.

Python
from collections import Counter

def find_duplicates_general(items: list) -> list:
    """Works for any hashable type."""
    counts = Counter(items)
    return [item for item, count in counts.items() if count > 1]

# Deduplicating LLM outputs
responses = ["answer A", "answer B", "answer A", "answer C", "answer B"]
print(find_duplicates_general(responses))   # ["answer A", "answer B"]

Problem 5: Rotate Array

Problem: Rotate an array to the right by k steps in-place.

Python
def rotate(nums: list[int], k: int) -> None:
    """
    Rotate array right by k positions in-place.
    O(n) time, O(1) space using the reverse trick.
    """
    n = len(nums)
    k = k % n   # Handle k larger than n

    def reverse(lo: int, hi: int) -> None:
        while lo < hi:
            nums[lo], nums[hi] = nums[hi], nums[lo]
            lo += 1
            hi -= 1

    # Step 1: reverse entire array
    reverse(0, n - 1)
    # Step 2: reverse first k elements
    reverse(0, k - 1)
    # Step 3: reverse remaining n-k elements
    reverse(k, n - 1)


arr = [1, 2, 3, 4, 5, 6, 7]
rotate(arr, 3)
print(arr)   # [5, 6, 7, 1, 2, 3, 4]

# Simpler: using slicing (creates a copy)
def rotate_slice(nums: list[int], k: int) -> list[int]:
    k = k % len(nums)
    return nums[-k:] + nums[:-k]

print(rotate_slice([1, 2, 3, 4, 5, 6, 7], 3))   # [5, 6, 7, 1, 2, 3, 4]

Problem 6: Flatten Nested Lists

Problem: Given a deeply nested list, return all elements as a flat list.

Python
from typing import Any

def flatten(nested: list[Any]) -> list[Any]:
    """
    Recursively flatten nested lists of any depth.
    
    AI context: flattening chunked document batches or ragged retrieval results.
    """
    result = []
    for item in nested:
        if isinstance(item, list):
            result.extend(flatten(item))
        else:
            result.append(item)
    return result


# Ragged retrieval results from multiple queries
multi_query_results = [
    ["doc1", "doc2"],
    ["doc3"],
    ["doc2", "doc4", "doc5"],
]

flat = flatten(multi_query_results)
print(flat)   # ["doc1", "doc2", "doc3", "doc2", "doc4", "doc5"]

# Deduplicated (preserving first-seen order)
seen = set()
unique_flat = [x for x in flat if not (x in seen or seen.add(x))]
print(unique_flat)   # ["doc1", "doc2", "doc3", "doc4", "doc5"]


# Iterative version (avoids recursion limit for very deep nesting)
def flatten_iterative(nested: list[Any]) -> list[Any]:
    result = []
    stack = [nested]

    while stack:
        current = stack.pop()
        if isinstance(current, list):
            stack.extend(reversed(current))   # reversed: left-to-right order
        else:
            result.append(current)

    return result

Problem 7: Product of Array Except Self

Problem: Given an array, return a new array where each element is the product of all other elements. No division allowed, O(n) time.

Python
def product_except_self(nums: list[int]) -> list[int]:
    """
    result[i] = product of all nums except nums[i].
    O(n) time, O(1) extra space (output array excluded).
    """
    n = len(nums)
    result = [1] * n

    # Forward pass: result[i] = product of all elements to the LEFT of i
    prefix = 1
    for i in range(n):
        result[i] = prefix
        prefix *= nums[i]

    # Backward pass: multiply by product of all elements to the RIGHT of i
    suffix = 1
    for i in range(n - 1, -1, -1):
        result[i] *= suffix
        suffix *= nums[i]

    return result


print(product_except_self([1, 2, 3, 4]))   # [24, 12, 8, 6]
# result[0] = 2*3*4 = 24
# result[1] = 1*3*4 = 12
# result[2] = 1*2*4 = 8
# result[3] = 1*2*3 = 6

Complexity Summary

| Problem | Time | Space | Key Technique | |---|---|---|---| | Two Sum | O(n) | O(n) | Hash map (seen dict) | | Sliding Window Max | O(n) | O(k) | Monotonic deque | | Merge Sorted Arrays | O(m+n) | O(m+n) | Two pointers | | Find Duplicates | O(n) | O(1) | Sign-flip trick | | Rotate Array | O(n) | O(1) | Triple reverse | | Flatten Nested | O(n) | O(depth) | Recursion or stack | | Product Except Self | O(n) | O(1) | Prefix + suffix pass |