"""Just using numpy.""" import numpy as np import sys def apply_gravity(terrain): tumbled = terrain.copy() while True: np.divmod(terrain, 4, tumbled, terrain) if not np.any(tumbled): return terrain[1:, :] += tumbled[:-1, :] terrain[:-1, :] += tumbled[1:, :] terrain[:, 1:] += tumbled[:, :-1] terrain[:, :-1] += tumbled[:, 1:] def main(height, show=True): width = int(height**0.5) + 1 terrain = np.zeros((width, width), dtype=np.int64) terrain[width // 2, width // 2] = height apply_gravity(terrain) if show: np.set_printoptions(threshold=sys.maxsize, linewidth=999) print(terrain) if __name__ == "__main__": main(int(sys.argv[1]))