Группировка объектов по заданной схеме, с помощью роевого алгоритма

Интересует доработка программы на языке Python. Необходимо реализовать группировку объектов по заданной схеме, в моём случае это квадрат. На данный момент моя программа вычисляет общий центр масс группы дронов, от чего те смещаются в общий центр. То, что должно получиться изображено на Изображении 1. Но есть нюанс,
на Изоражении 1, было создано 100 дронов, и квадрат формируется полноценно. Но необходимо учесть возможность, если бы дронов не хватало, допустим было бы 97, хотелось бы, чтобы дроны точно так же формировали квадрат, смещаясь в общий центр, но при этом недостроенный (Пример на Изображении 2). Вся программа в архиве.
Так же необходимо составить документ, какие изменения были внесены, и что было улучшено.