Advertisement
hhoppe

Advent of code 2024 day 20 numpy fully vectorized

Dec 21st, 2024
35
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 1.58 KB | None | 0 0
  1. def day20(s, *, part2=False, min_savings=100):  # Fully vectorized numpy solution.  Not fastest.
  2.   grid = np.array([list(line) for line in s.splitlines()])
  3.   ((ys, xs),) = np.argwhere(grid == 'S')
  4.   ((ye, xe),) = np.argwhere(grid == 'E')
  5.   distance = np.where(grid == '#', -2, -1)
  6.   neighbors = (0, -1), (0, 1), (-1, 0), (1, 0)
  7.  
  8.   d, y, x = 0, ys, xs
  9.   distance[y, x] = d
  10.   while (y, x) != (ye, xe):
  11.     ((y, x),) = (yx2 for dy, dx in neighbors if distance[yx2 := (y + dy, x + dx)] == -1)
  12.     distance[y, x] = d = d + 1
  13.  
  14.   radius = 20 if part2 else 2
  15.   y, x = np.indices((radius + 1, 2 * radius + 1))
  16.   x -= radius
  17.   valid = (abs(y) + abs(x) <= radius) & ((y > 0) | (y == 0) & (x > 0))
  18.   shifts = np.argwhere(valid)
  19.   y_shifts = shifts[:, 0]
  20.   x_shifts = shifts[:, 1] - radius
  21.   manhattan = abs(y_shifts) + abs(x_shifts)
  22.  
  23.   padded = np.pad(distance, radius, constant_values=-1)
  24.  
  25.   # Create a view for vectorized shifting.
  26.   ny, nx = distance.shape
  27.   shifted_views = padded[
  28.       radius + y_shifts[:, None, None] + np.arange(ny)[None, :, None],
  29.       radius + x_shifts[:, None, None] + np.arange(nx)[None, None, :],
  30.   ]
  31.  
  32.   # Create mask for valid (non-negative) values.
  33.   valid_mask = distance >= 0
  34.   shifted_valid = shifted_views >= 0
  35.  
  36.   # Calculate value differences and adjust by Manhattan distance.
  37.   value_diff = np.abs(shifted_views - distance)
  38.   adjusted_diff = value_diff - manhattan[:, None, None]
  39.  
  40.   # Count pairs satisfying all conditions.
  41.   valid_pairs = ((adjusted_diff >= min_savings) & valid_mask & shifted_valid).sum(axis=(1, 2))
  42.   return valid_pairs.sum()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement