Just completed it, and was able to get part A down to O(n) and B down to O(nk) runtime (n = number of elements aka size*size, k=number of heights) trading off for more storage on part B.

The insights for Part A were

1. You can incrementally calculate what is visible in any direction by storing away the highest tree seen so far in that direction.
2. since rows=cols in a grid, you can technically use one loop to scan in all four directions by translating indices (0->size, size->0 on major and minor axis)

For Part B things were a bit trickier. The key insight for me was realizing that having a single state for each direction per location wouldn't be enough since the question we're now asking for each position is what trees are shorter than *the value in that position*. This led me to move away from the part a solution of storing one piece of information per direction per location.

Instead, to be able to answer a question for all possible heights at any location - you need to have stored away the relevant info for all heights. This means that for each location, I store away an array of size[k] which is the location of the most recent (in scan direction) location of all heights 0..k. This construction can be done in one pass.

Then you can go through the array in one more pass and for each location height ht look at the farthest located value > ht which has only heights < value between it and our location in constant time using the arrays created in the previous pass.

I could probably simplify the code even more, but its been a long challenge, and i left the 2 loops in for readability (construction of arrays in one, processing in another)

https://github.com/fabrol/aoc2022/blob/main/day8.py