Spaces:
Running
Running
File size: 3,108 Bytes
a4da721 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 |
def load_data(file):
with open(file) as f:
raw_data = f.readlines()
grid = []
for line in raw_data:
line = line.strip("\n")
grid.append(list(line))
return grid
def get_antennas(grid):
M = len(grid)
N = len(grid[0])
antennas = {}
for i in range(M):
for j in range(N):
if grid[i][j] != ".":
a = grid[i][j]
if antennas.get(a):
antennas[a].append((i,j))
else:
antennas[a] = [(i,j)]
return antennas
def get_next_nodes(a, b):
# Get direction vector
dx = b[0] - a[0]
dy = b[1] - a[1]
node_a = (a[0] - dx, a[1] - dy) # Subtract from first node
node_b = (b[0] + dx, b[1] + dy) # Add to second node
return node_a, node_b
def add_antinode(a, grid):
M = len(grid)
N = len(grid[0])
i,j = a
if i in range(M) and j in range(N):
grid[i][j] = '#'
def count_antinodes(grid):
M = len(grid)
N = len(grid[0])
total = 0
for i in range(M):
for j in range(N):
if grid[i][j] == '#':
total += 1
return total
def add_antinodes(a, b, grid):
next_nodes = get_next_nodes(a,b)
for node in next_nodes:
add_antinode(node, grid)
grid = load_data("input.txt")
antennas = get_antennas(grid)
for freq in antennas:
# Get all positions for a given freq.
positions = antennas[freq]
for i in range(len(positions)):
for j in range(i+1, len(positions)):
# For each pair of antennas, add the antinodes
a, b = positions[i], positions[j]
add_antinodes(a, b, grid)
print(count_antinodes(grid))
## Part 2
def get_direction_vector(a, b):
# Get direction vector
dx = b[0] - a[0]
dy = b[1] - a[1]
return dx, dy
def is_in_bounds(node, grid):
M = len(grid)
N = len(grid[0])
i, j = node
return i in range(M) and j in range(M)
def add_antinodes(a, b, grid):
dx, dy = get_direction_vector(a,b)
next_node = (a[0] - dx, a[1] - dy) # Subtract from first node
while is_in_bounds(next_node, grid):
add_antinode(next_node, grid)
next_node = (next_node[0] - dx, next_node[1] - dy)
next_node = (b[0] + dx, b[1] + dy) # Add to second node
while is_in_bounds(next_node, grid):
add_antinode(next_node, grid)
next_node = (next_node[0] + dx, next_node[1] + dy)
def count_antinodes(grid):
M = len(grid)
N = len(grid[0])
total = 0
for i in range(M):
for j in range(N):
if grid[i][j] != '.':
# Antennas count too
total += 1
return total
grid = load_data("input.txt")
antennas = get_antennas(grid)
for freq in antennas:
positions = antennas[freq]
for i in range(len(positions)):
for j in range(i+1, len(positions)):
a, b = positions[i], positions[j]
add_antinodes(a, b, grid)
print(count_antinodes(grid))
def pprint(grid):
grid_str = "\n".join(["".join(l) for l in grid])
print(grid_str)
|