Hopper vs Go-to-the-middle

main.py

import copy
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

fig, ax = plt.subplots()
ax.set_title("(iii) hopper")
xdata, ydata = [], []
ln, = plt.plot([], [], 'ro', animated=True)
plt.axis([0, 6, 0, 3])

def hopper(robots):
    idx = 0
    N = len(robots)
    results = [copy.deepcopy(robots)]
    remove = [0]
    robots_ = copy.deepcopy(robots)
    a_ = None

    while True:
        remove = 0; idx = 1
        while idx > 0 and idx < N - 1 - remove:
            i = idx - remove
            if dist(robots_[i - 1], robots[i + 1]) <= 1:
                robots_ = np.delete(robots_, i, 0)
                remove += 1
                break
            elif deg(robots_[i - 1], robots_[i], robots_[i + 1]) <= 90:
                robots_[i][0] = (robots_[i - 1][0] + robots_[i + 1][0]) / 2
                robots_[i][1] = (robots_[i - 1][1] + robots_[i + 1][1]) / 2
                break
            else:
                if i > 0 and i < N - 1 - remove:
                    v = robots_[i+1] - robots_[i]
                    robots_[i] = robots_[i-1] + v
                elif i == 0:
                    v_ = robots_[i+1] - robots_[i]
                elif i == N - 1 - remove:
                    robots_[i] = robots_[i-1] + v_
            idx += 1
        results.append(robots_)
        
        if np.allclose(expected_result(robots_), robots_, atol=1e-4):
            break
        
        robots_ = copy.deepcopy(robots_)

    return results

def go_to_the_middle(robots):
    N = len(robots)
    results = [copy.deepcopy(robots)]
    result = copy.deepcopy(robots)
    expected_res = expected_result(robots)

    idx = 0
    while True:
        for i in range(1, N - 1):
            if deg(robots[i - 1], robots[i], robots[i + 1]) != 180:
                result[i][0] = (robots[i - 1][0] + robots[i + 1][0]) / 2
                result[i][1] = (robots[i - 1][1] + robots[i + 1][1]) / 2
        results.append(copy.deepcopy(result))
        idx += 1
        if np.allclose(expected_res, result, atol=1e-4):
            break
        robots = copy.deepcopy(result)
        
    return results

def dist(a, b):
    return np.linalg.norm(a - b)

def deg(a, b, c):
    ba = a - b
    bc = c - b
    angle = np.arccos(np.dot(ba, bc) / (np.linalg.norm(ba) * np.linalg.norm(bc)))
    return np.degrees(angle)

def update_robots(robots):
    ln.set_data(*zip(*robots))
    return ln,

def expected_result(robots):
    expected = np.zeros_like(robots)
    N = len(robots)
    expected[0] = robots[0]
    expected[-1] = robots[-1]
    for i in range(1, N-1):
        expected[i] = (expected[-1] - expected[0]) / (N-1) * i + expected[0]
    return expected

def main():
    ###### (i + ii) case: hopper bettern than go-to-the-middle ######
    robots = np.array([[1, 1], [2, 2], [3, 3], [4, 2], [5, 1], [6, 0], [7, 1]], dtype='float64')
    robots = np.array([[0.1, 1], [0.2, 2], [0.3, 1], [0.4, 2], [0.5, 1]])    
    ###### (iii) case: hopper worse than go-to-the-middle #####
    robots = np.array([[1, 1], [2, 2], [3, 1], [4, 0], [5, 1]], dtype='float64')
    
    ##### Select method ######
    results = hopper(robots)
    results = go_to_the_middle(robots)

    ###### Save GIF file #####
    ani = FuncAnimation(fig, update_robots, results, blit=True)
    ani.save('result-iii-hopper.gif', dpi=80, writer='imagemagick')

if __name__ == '__main__':
    main()