Python – Find the Union and Intersection of Two Sorted Arrays

This tutorial explains how to find the union and intersection of two sorted arrays in Python. These operations are common in data structure and algorithm interviews. We will break down the problem statement, sample input/output, solution steps, and provide a complete Python program.

Problem Statement

Given two sorted arrays, you are required to:

  1. Find the union: a sorted list containing all unique elements from both arrays.
  2. Find the intersection: a sorted list of elements common to both arrays.

Sample Input and Output

Example 1:

</>
Copy
Input: arr1 = [1, 3, 4, 5, 7]
       arr2 = [2, 3, 5, 6]

Output: Union: [1, 2, 3, 4, 5, 6, 7]
        Intersection: [3, 5]

Example 2:

</>
Copy
Input: arr1 = [0, 1, 2, 3]
       arr2 = [2, 3, 4, 5]

Output: Union: [0, 1, 2, 3, 4, 5]
        Intersection: [2, 3]

Solution Approach

Since the arrays are already sorted, we can use a two-pointer approach to efficiently find both the union and intersection:

  1. Initialize two pointers, one for each array.
  2. Compare the elements at the current pointers:
    • If arr1[i] is less than arr2[j], add arr1[i] to the union (if not already added) and move the pointer in arr1 forward.
    • If arr2[j] is less than arr1[i], add arr2[j] to the union (if not already added) and move the pointer in arr2 forward.
    • If both elements are equal, add the element to both the union and intersection (avoiding duplicates), and move both pointers forward.
  3. Once one of the arrays is fully traversed, add the remaining elements from the other array to the union.

Python Program

The following Python program implements the above approach using a two-pointer technique:

</>
Copy
def union_intersection(arr1, arr2):
    n1, n2 = len(arr1), len(arr2)
    i, j = 0, 0
    union_result = []
    intersection_result = []

    # Traverse both arrays simultaneously
    while i < n1 and j < n2:
        # If current element of arr1 is smaller, add it to union
        if arr1[i] < arr2[j]:
            if not union_result or union_result[-1] != arr1[i]:
                union_result.append(arr1[i])
            i += 1
        # If current element of arr2 is smaller, add it to union
        elif arr1[i] > arr2[j]:
            if not union_result or union_result[-1] != arr2[j]:
                union_result.append(arr2[j])
            j += 1
        else:
            # Both elements are equal, add to union and intersection
            if not union_result or union_result[-1] != arr1[i]:
                union_result.append(arr1[i])
            intersection_result.append(arr1[i])
            i += 1
            j += 1

    # Add any remaining elements from arr1
    while i < n1:
        if not union_result or union_result[-1] != arr1[i]:
            union_result.append(arr1[i])
        i += 1

    # Add any remaining elements from arr2
    while j < n2:
        if not union_result or union_result[-1] != arr2[j]:
            union_result.append(arr2[j])
        j += 1

    return union_result, intersection_result

# Example 1
arr1 = [1, 3, 4, 5, 7]
arr2 = [2, 3, 5, 6]
union_res, intersection_res = union_intersection(arr1, arr2)
print("Example 1:")
print("Union:", union_res)
print("Intersection:", intersection_res)

# Example 2
arr1 = [0, 1, 2, 3]
arr2 = [2, 3, 4, 5]
union_res, intersection_res = union_intersection(arr1, arr2)
print("\nExample 2:")
print("Union:", union_res)
print("Intersection:", intersection_res)

Another variation of the program, with union and intersection operations written separately.

</>
Copy
def add_element_to_result(result, element):
    if not result or result and result[-1] != element:
        result.append(element)

def union(arr1, arr2):
    p1 = p2 = 0

    union_result = []

    while p1 < len(arr1) and p2 < len(arr2):
        if arr1[p1] < arr2[p2]:
            add_element_to_result(union_result, arr1[p1])
            p1 += 1
        elif arr2[p2] < arr1[p1]:
            add_element_to_result(union_result, arr2[p2])
            p2 += 1
        else:
            add_element_to_result(union_result, arr2[p2])
            p1 += 1
            p2 += 1

    for item in arr1[p1:]:
        if union_result and union_result[-1] != arr1[p1]:
            union_result.append(arr1[p1])

    for item in arr2[p2:]:
        if union_result and union_result[-1] != arr2[p2]:
            union_result.append(arr2[p2])

    return union_result


def intersection(arr1, arr2):
    p1 = p2 = 0

    intersection_result = []

    while p1 < len(arr1) and p2 < len(arr2):
        if arr1[p1] < arr2[p2]:
            p1 += 1
        elif arr2[p2] < arr1[p1]:
            p2 += 1
        else: #arr1[p1] == arr2[p2]
            add_element_to_result(intersection_result, arr2[p2])
            p1 += 1
            p2 += 1

    return intersection_result


arr1 = [1, 3, 4, 5, 7, 9, 10, 11, 12]
arr2 = [2, 3, 5, 6, 15]
union_result = union(arr1, arr2)
intersection_result = intersection(arr1, arr2)
print('Union :', union_result)
print('Intersection :', intersection_result)

Conclusion

In this tutorial, we demonstrated how to find the union and intersection of two sorted arrays using a two-pointer approach. This method is efficient with a time complexity of O(n + m), where n and m are the lengths of the two arrays.