visualizer.py

import matplotlib.pyplot as plt
import matplotlib.animation as animation
import random
import argparse
from typing import List, Generator, Tuple


def bubble_sort(arr: List[int]) -> Generator[Tuple[List[int], int, int], None, None]:
    n = len(arr)
    for i in range(n):
        for j in range(n - i - 1):
            if arr[j] > arr[j + 1]:
                arr[j], arr[j + 1] = arr[j + 1], arr[j]
            yield arr, j, j + 1


def selection_sort(arr: List[int]) -> Generator[Tuple[List[int], int, int], None, None]:
    n = len(arr)
    for i in range(n):
        min_idx = i
        for j in range(i + 1, n):
            if arr[j] < arr[min_idx]:
                min_idx = j
            yield arr, i, j
        arr[i], arr[min_idx] = arr[min_idx], arr[i]
        yield arr, i, min_idx


def insertion_sort(arr: List[int]) -> Generator[Tuple[List[int], int, int], None, None]:
    for i in range(1, len(arr)):
        key = arr[i]
        j = i - 1
        while j >= 0 and arr[j] > key:
            arr[j + 1] = arr[j]
            j -= 1
            yield arr, j + 1, i
        arr[j + 1] = key
        yield arr, j + 1, i


def quick_sort(arr: List[int], low: int = 0, high: int = None) -> Generator[Tuple[List[int], int, int], None, None]:
    if high is None:
        high = len(arr) - 1
    if low < high:
        pivot_idx = yield from partition(arr, low, high)
        yield from quick_sort(arr, low, pivot_idx - 1)
        yield from quick_sort(arr, pivot_idx + 1, high)


def partition(arr: List[int], low: int, high: int) -> Generator[Tuple[List[int], int, int], None, int]:
    pivot = arr[high]
    i = low - 1
    for j in range(low, high):
        if arr[j] <= pivot:
            i += 1
            arr[i], arr[j] = arr[j], arr[i]
        yield arr, i, j
    arr[i + 1], arr[high] = arr[high], arr[i + 1]
    yield arr, i + 1, high
    return i + 1


def merge_sort(arr: List[int], left: int = 0, right: int = None) -> Generator[Tuple[List[int], int, int], None, None]:
    if right is None:
        right = len(arr) - 1
    if left < right:
        mid = (left + right) // 2
        yield from merge_sort(arr, left, mid)
        yield from merge_sort(arr, mid + 1, right)
        yield from merge(arr, left, mid, right)


def merge(arr: List[int], left: int, mid: int, right: int) -> Generator[Tuple[List[int], int, int], None, None]:
    left_copy = arr[left:mid + 1]
    right_copy = arr[mid + 1:right + 1]
    i = j = 0
    k = left
    while i < len(left_copy) and j < len(right_copy):
        if left_copy[i] <= right_copy[j]:
            arr[k] = left_copy[i]
            i += 1
        else:
            arr[k] = right_copy[j]
            j += 1
        k += 1
        yield arr, k, k - 1
    while i < len(left_copy):
        arr[k] = left_copy[i]
        i += 1
        k += 1
        yield arr, k, k - 1
    while j < len(right_copy):
        arr[k] = right_copy[j]
        j += 1
        k += 1
        yield arr, k, k - 1


ALGORITHMS = {
    "bubble": bubble_sort,
    "selection": selection_sort,
    "insertion": insertion_sort,
    "quick": quick_sort,
    "merge": merge_sort,
}


def animate(algorithm_name: str, size: int = 50, interval: int = 10):
    data = [random.randint(1, 100) for _ in range(size)]
    generator = ALGORITHMS[algorithm_name](data)

    fig, ax = plt.subplots(figsize=(12, 6))
    ax.set_title(f"{algorithm_name.title()} Sort")
    bar_container = ax.bar(range(len(data)), data, color="steelblue")
    ax.set_xlim(-1, len(data))
    steps_done = [0]

    def update(_):
        try:
            arr, i, j = next(generator)
            for bar, val in zip(bar_container, arr):
                bar.set_height(val)
            for idx, bar in enumerate(bar_container):
                if idx in (i, j):
                    bar.set_color("crimson")
                elif idx <= steps_done[0] and algorithm_name in ("bubble", "selection", "insertion"):
                    bar.set_color("limegreen")
                else:
                    bar.set_color("steelblue")
            steps_done[0] += 1
        except StopIteration:
            for bar in bar_container:
                bar.set_color("limegreen")

        return bar_container

    ani = animation.FuncAnimation(fig, update, frames=range(1000), interval=interval, blit=False, repeat=False)
    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Visualize sorting algorithms")
    parser.add_argument("algorithm", choices=list(ALGORITHMS.keys()), help="Sorting algorithm")
    parser.add_argument("--size", type=int, default=50, help="Number of elements to sort")
    parser.add_argument("--interval", type=int, default=10, help="Animation speed (ms per frame)")
    args = parser.parse_args()
    animate(args.algorithm, args.size, args.interval)