fast-abelian-sandpile/sandpile_numpy.py

31 lines
746 B
Python

"""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]))