My initial solution essentially used brute force to check to the left, right, up and down of each element to see if there was an element greater than or equal to the current element to determine if the current tree is visible from outside the forest. I used short-circuit evaluation to optimize this solution, but it’s still $O(n*\sqrt{n})$ since we examine $O(\sqrt{n})$ trees for each of the $n$ trees in the forest to determine their visibility. This is because many of the trees are checked multiple times during the brute force solution. (note: $n$ here is the number of trees, not the number of rows of trees. Thus, assuming the forest is square and has width $k$, then the forest has $n = k^2$ trees and $\sqrt{n}$ trees per row/column )

To arrive at an $O(n)$ solution, I applied the next greater element algorithm that uses a monotonically decreasing stack to find the value (or index) of the next greater element to the right (or left) of each element in an input array in linear time.

next greater element

To find the next greater element to the right of each element in an array, we will add the indices of elements from the input array to a stack until we find an element that’s larger than the top element on the stack. This means the stack is monotonically decreasing from bottom to top, since we only push an element onto the stack if it’s smaller than the top element. Once we find an element larger than the top element, this means some of the items in the stack are smaller than this element, meaning this element is their next greater element (to the right). We then pop the top index off the stack and set the answer array at that index to the current element. We continue doing this until the element corresponding to the index at the top of the stack is no longer smaller than the current element, or until the stack is empty. Then, we continue iterating through the array, repeating the above steps until we’ve added all elements to the stack. The indices left over in the stack indicate the values of the array that did not have any greater elements to their right. To indicate this, the value at these indices in the answer array should be $-1$.

Here is the code to help this make more sense:

def next_greater_right(arr):
    res = [-1 for i in arr] # start all indices at -1
    stack = []
    for i, num in enumerate(arr):
        while len(stack) > 0 and arr[stack[-1]] < num:
            index = stack.pop()
            res[index] = num
        stack.append(i)
    return res

Notice the first time the while loop is encountered, it will not run since the stack is empty. Here’s an example output of this algorithm:

arr: [1, 2, 5, 13, 9, 6, 4, 12, 7]
answer: [2, 5, 13, -1, 12, 12, 12, -1, -1]

To find the next greater element to the right of an element at index $i$ in $arr$, you just need to access $answer[i]$. Notice how $answer[3] = -1$ since there are no elements greater than 13 to the right of 13.

To find the greater element to the left of each element in an array, simply reverse the order of iteration of the for loop so you’re comparing elements with other elements to their left:

def next_greater_left(arr):
    res = [-1 for i in arr]
    stack = []
    for i, num in reversed(list(enumerate(arr))): # iterating from left to right
        while len(stack) > 0 and arr[stack[-1]] < num:
            index = stack.pop()
            res[index] = num
        stack.append(i)
    return res

applying next greater element to tree visibility

To determine if a tree is visible from the outside, we need to know if it has a greater or equal element to its left, right, up or down directions. If it does not have a greater or equal element for one or more of these directions, then it is visible from outside the forest. Thus, to determine whether each tree is visible from the right, we need to run the next_greater_right function on each row of the forest. To determine if each tree is visible from the left, we run next_greater_left on each row of the forest. To determine if each tree is visible from below, we can transpose the forest (turning columns into rows) and then run next_greater_right on each row of the transpose (iterating from the top down in the forest’s original configuration), then transpose the result, reverting the answer back to the original forest structure. Similarly, to determine if each tree is visible from above, first transpose the forest and then run next_greater_left on each row of the transpose (iterating from bottom up in the original configuration) then transpose the result, reverting the answer back to its original structure:

def transpose(mat): # transposes a rectangular 2d array
    res = []
    w = len(mat[0])
    h = len(mat)
    for j in range(w):
        res.append([mat[i][j] for i in range(h)])
    return res

def part1():
    with open('input.txt') as f:
        # turn lines of input into 2d array of integer heights
        forest = list(map(lambda l: l.strip(), f.readlines()))
        forest = [[int(c) for c in row] for row in forest]
        transposed_forest = transpose(forest)

        blocking_trees_right = [next_greater_right(row) for row in forest]
        blocking_trees_left = [next_greater_left(row) for row in forest]
        blocking_trees_above = transpose([next_greater_left(row) for row in transposed_forest])
        blocking_trees_below = transpose([next_greater_right(row) for row in transposed_forest])

Now we have four grids where each element of a particular grid represents the height of the tree blocking that element in the specified direction OR, if the value is $-1$, it indicates that there are no trees blocking the current element in the specified direction. Next, we need to iterate through each tree in the forest and check the corresponding element in these four grids. If any of the values are -1, then this means there’s a clear line of sight from that tree to an edge of the forest meaning it’s visible. Otherwise, the tree is blocked on all sides:

visible_trees = {}
        for i, row in enumerate(forest):
            for j, col in enumerate(forest[i]):
                if (blocking_trees_left[i][j] == -1 
                    or blocking_trees_right[i][j] == -1 
                    or blocking_trees_above[i][j] == -1 
                    or blocking_trees_below[i][j] == -1
                    ):
                    visible_trees[f'{i},{j}'] = forest[i][j]

Below is the full solution for part 1: