def distance2(a, b, N):
    """Compute the distance between two vectors with wrapping of size N."""
    d2 = 0
    for i in range(a.shape[0]):
        aa = a[i]
        bb = b[i]
        dx = aa - bb
        if dx < -N/2:
            d2 += (dx + N)**2
        elif dx > N/2:
            d2 += (dx - N)**2
        else:
            d2 += (dx)**2
    return d2