parallel_merge.py

from random import shuffle
from math import log2, floor

def main():
    A = list(range(15))
    B = list(range(10))
    print("Before:\t A =", A)
    print("\t B =", B)
    N = len(A)
    P = 5 # Number of procs

    block_size = floor(N / P) # minimum block size

    b_block_start = 0

    result = []
    for p in range(P):
        a_block_start = p * block_size + min(p, N % P)
        a_block_size = block_size + (1 if p < N % P else 0)

        last_val_in_a_block = A[a_block_start + a_block_size - 1]
        b_block_end = last_index_less_or_eq_to(B,
                                               last_val_in_a_block,
                                               b_block_start)
        if b_block_end is None:
            b_block_size = 0
        else:
            b_block_size = b_block_end - b_block_start + 1
        merge_result =  merge(A, a_block_start, a_block_size,
                              B, b_block_start, b_block_size)
        result += merge_result
        b_block_start += b_block_size
    print("Merged:\t", result)
    
def last_index_less_or_eq_to(xs, value, start_index):
    prev_index = None
    for i, x in enumerate(xs[start_index:]):
        if x > value:
            break
        else:
            prev_index = i + start_index
    return prev_index 
    
def merge(A, a_start, a_size, B, b_start, b_size):
    result = []
    a_index = a_start
    b_index = b_start
    
    # because of the way we choose b_size, the last value in the B block
    # will always be < last value in A block, so B block will be consumed
    # first
    while b_index < b_start + b_size:
        if A[a_index] < B[b_index]:
            result.append(A[a_index])
            a_index += 1
        else:
            result.append(B[b_index])
            b_index += 1

    while a_index < a_start + a_size:
        result.append(A[a_index])
        a_index += 1
    return result

main()