Bitonic sequences

bitonic.py

import numpy as np
from tqdm.auto import tqdm

def is_bitonic(xs):
    results = []
    for seq in xs:
        prev = seq[0]
        direction = 0
        num_changes = 0
        for x1 in seq:
            if x1 > prev:
                if direction == -1:
                    direction = 1
                    num_changes += 1
                elif direction == 0:
                    direction = 1
            elif x1 < prev:
                if direction == 1:
                    direction = -1
                    num_changes += 1
                elif direction == 0:
                    direction = -1
            prev = x1
        results.append(num_changes < 2)
    return np.array(results)

def is_tritonic(xs):
    results = []
    for seq in xs:
        prev = seq[0]
        direction = 0
        num_changes = 0
        for x1 in seq:
            if x1 > prev:
                if direction == -1:
                    direction = 1
                    num_changes += 1
                elif direction == 0:
                    direction = 1
            elif x1 < prev:
                if direction == 1:
                    direction = -1
                    num_changes += 1
                elif direction == 0:
                    direction = -1
            prev = x1
        results.append(num_changes < 3)
    return np.array(results)

assert np.all(is_bitonic(np.array([
    [0, 1, -1, 2],
    [0, 1, -1, -2],
    [0, 1, -1, -2],
    [0, 0, 1, -1],
    [0, -1, -1, 0],
    [0, 0, 1, 1],
])) == np.array([False, True, True, True, True, True]))
assert np.all(is_bitonic(np.array([[0, 1, 2, 3, 4, 5]])) == np.array([True]))
assert np.all(is_bitonic(np.array([[0, 1, 1, 1, 2, 2, 2, 3, 3]])) == np.array([True]))
assert np.all(is_bitonic(np.array([[0, 1, 1, 1, 2, 2, 2, 3, 3, 2, 2, 1, 1]])) == np.array([True]))
assert np.all(is_bitonic(np.array([[0, 1, 1, 1, 2, 2, 2, 3, 3, 2, 2, 1, 1, 2]])) == np.array([False]))

for i in tqdm(range(100)):
    n = 16384
    items = np.random.normal(size=[n])
    expected = np.sort(items)
    items = np.sort(items.reshape([-1, n // 2]), axis=-1).ravel()                

    # Bitonic sort step 1, with reversal
    for i in range(n // 2):
        i1 = len(items) - i - 1
        x = items[i]
        y = items[i1]
        if x > y:
            items[i], items[i1] = y, x

    # Step 2
    stride = n // 4
    while stride >= 1:
        if stride > 1:
            assert np.all(is_tritonic(items.reshape([-1, 2*stride])))
        for i in range(len(items)):
            if i & stride:
                continue
            i1 = i | stride
            x, y = items[i], items[i1]
            if x > y:
                items[i], items[i1] = y, x
        stride //= 2
    
    assert np.allclose(items, expected)