Rust

Using a HashMap for the octopuses instead of a 2D array worked pretty well together with helper functions from my Coordinates class.

Stepping logic for both parts

fn step(&mut self) -> usize {
    // First, the energy level of each octopus increases by 1.
    for c in Coordinates::in_area(0..self.width, 0..self.height) {
        self.octopuses.entry(c).and_modify(|e| *e += 1);
    }
    let mut flashed = HashSet::new();
    loop {
        // Then, any octopus with an energy level greater than 9 flashes.
        let flashes = Coordinates::in_area(0..self.width, 0..self.height)
            .filter(|c| *self.octopuses.get(c).unwrap() > 9)
            .filter(|c| !flashed.contains(c))
            .collect::<HashSet<Coordinates>>();

        // This increases the energy level of all adjacent octopuses by 1,
        // including octopuses that are diagonally adjacent.
        for c in flashes.iter().flat_map(|c| c.all_offset_by(1)) {
            self.octopuses.entry(c).and_modify(|e| *e += 1);
        }

        // If this causes an octopus to have an energy level greater than 9, it also flashes.
        // This process continues as long as new octopuses keep having their energy level
        // increased beyond 9.
        if flashes.is_empty() {
            break;
        }
        flashed.extend(flashes);
    }

    // Finally, any octopus that flashed during this step has its energy level set to 0,
    // as it used all of its energy to flash.
    for c in flashed.iter() {
        self.octopuses.entry(*c).and_modify(|e| *e = 0);
    }

    // How many total flashes are there after this step?
    flashed.len()
}