You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
32 lines
819 B
32 lines
819 B
3 years ago
|
require "set"
|
||
|
|
||
|
INPUT = ARGF.read.split("\n").map {|l| l.chars.map(&:to_i) }
|
||
|
|
||
|
NEIGHBORS = [[0,-1], [-1, 0], [1, 0], [0, 1]]
|
||
|
def neighbors(y,x)
|
||
|
NEIGHBORS
|
||
|
.map {|dy,dx| [y+dy, x+dx] }
|
||
|
.select {|y,x| (0...INPUT.size).cover?(y) && (0...INPUT[0].size).cover?(x) }
|
||
|
end
|
||
|
|
||
|
all = (0...INPUT.size).flat_map {|y| (0...INPUT[0].size).map {|x| [y, x]}}
|
||
|
lows = all.select {|y,x| neighbors(y,x).all? {|yy,xx| INPUT[yy][xx] > INPUT[y][x] }}
|
||
|
# p lows.sum {|y,x| INPUT[y][x] + 1 }
|
||
|
|
||
|
def fill_basin(low)
|
||
|
basin = Set.new([low])
|
||
|
|
||
|
stack = [low]
|
||
|
until stack.empty?
|
||
|
y,x = stack.shift
|
||
|
new_neighbors = neighbors(y,x).reject {|y,x| basin.include?([y,x]) || INPUT[y][x] == 9 }
|
||
|
stack.concat(new_neighbors)
|
||
|
basin.merge(new_neighbors)
|
||
|
end
|
||
|
|
||
|
basin
|
||
|
end
|
||
|
|
||
|
p lows.map { fill_basin(_1) }.map(&:size).sort[-3, 3].inject(:*)
|
||
|
|