Egg balance

egg_balance.py

import numpy as np

rows = 3
cols = 6

# Create all 2^(rows*cols) combinations of egg matrices.
all_combos = np.zeros([1] * (rows * cols + 2), dtype=np.float32)
for i in range(rows * cols):
    m = np.zeros([2, rows * cols], dtype=np.float32)
    m[1, i] = 1
    m = m.reshape([1 if i != j else 2 for j in range(rows * cols)] + [rows, cols])
    all_combos = all_combos + m
all_combos = all_combos.reshape([-1, rows, cols])

# Compute the center of masses
coords = np.array([[row, col] for row in range(rows) for col in range(cols)])
unnormalized = np.sum(all_combos.reshape([-1, rows * cols, 1]) * coords, axis=1)
normalized = unnormalized / np.maximum(1e-5, np.sum(all_combos, axis=(1, 2))[:, None])

center_dists = np.sum((normalized - [1, 2.5]) ** 2, axis=-1)
is_center = center_dists < 1e-5
is_odd = np.round(np.sum(all_combos, axis=(1, 2))).astype(np.int32) % 2 == 1
print("total centered configs:", np.sum(is_center))
print("total odd configs:", np.sum(is_odd))
print("total odd centered configs:", np.sum(np.logical_and(is_center, is_odd)))

min_odd_idx = np.argmin(center_dists[is_odd])
min_odd = all_combos[is_odd][min_odd_idx]
print("most centered odd configuration:")
print(min_odd)