Dynamic programming and lazy computing

Dynamic programming is a rather important approach in solving many complex problems, based on splitting a task into subtasks and solving these subtasks once, even if they are part of several other subtasks. For people who are just beginning to master functional programming, the question often arises: “ How to avoid re-solving subtasks if you do not use variables to save the results? ". In this question, one feature of functional programming — the absence of variables — makes it difficult to cache the results of solving subtasks, but another feature will help — lazy calculations.

Consider the problem of finding watersheds:

Given a map of the area, represented by a rectangular matrix of heights (for example, integers). From each cell, water will drain into one of the neighboring cells (north, east, south, west) with the lowest height if
it is less than the height of this cell. Otherwise, this cell is called a drain. It is necessary to mark equally all the cells of the map, the water from which will flow into the same drain.

Thus, to find the drain of a single cell, you need to find the drain of its neighbor (if it is not a drain itself), and it is clear that water from several different cells can collect and flow along the same route into one drain.
A rather simple recursive algorithm comes to mind:

 drainOf cell = 
   case lowestNeighbourOf cell of
     Nothing → cell
     Just nb → drainOf nb

But in this case, the drain of each cell will be calculated many times, which can not but depress.
In this case, lazy calculations will help us: their essence is that the expression will not be calculated until its result is needed elsewhere, and after the calculation the result will be reused if necessary.
That's what we need!
The question arises: if in Haskell all (almost) calculations are lazy, then why won't the above code cache intermediate results?
The thing is that the following expression code is different and the function will be called twice!

 print (drainOf cell)
 print (drainOf cell)

But in the following code, the function will be called only once:

 let drain = drainOf cell
 print drain
 print drain

Those. we need to somehow indicate that we mean the same expressions, then lazy calculations will work for us.

 cells = ((1,1), (width, height))

 altitude :: (int, int) → int

 - here it is, an array of drains of all cells!  But it is calculated lazily
 - i.e.  for now this is only a list of calls to the drainOf function, but not its results.
 drains = array cells [(cell, drainOf cell) |  cell ← range cells]
  
 drainOf cell = 
   case lowestNeighbourOf cell of
     Nothing → cell
     Just nb → drains!  nb - appeal to the element of the array of stocks
     - here at the first call, the calculation of drainOf will occur,
     - and at the second - the return of the cached result

 lowestNeighbourOf (x, y) =
   let neighbors = filter (inRange cells) [ 
                         (x, y-1), (x + 1, y), (x, y + 1), (x-1, y)]
       lowest = minimumBy (compare`on`altitude) neighbors
   in
       if altitude lowest ≥ altitude (x, y)
             then Nothing
             else just lowest

That's all, I hope you have become clearer what lazy calculations are.

UPD: I decided to add the full program and its output, so as not to be unfounded.

 import Data.Array
 import Debug.Trace (trace)
 import Control.Monad (forM_)
 import Data.List (minimumBy, transpose)
 import Data.Function (on)

 - simple data -

 cells = ((1,1), (3,3))

 altitudes = listArray cells $ concat $ transpose [
     [10, 8, 10],
     [15, 3, 15],
     [20, 25, 20]]

 - common functions -

 altitude :: (int, int) → int
 altitude cell = altitudes!  cell

 lowestNeighbourOf (x, y) =
   let neighbors = filter (inRange cells) [ 
                         (x, y-1), (x + 1, y), (x, y + 1), (x-1, y)]
       lowest = minimumBy (compare`on`altitude) neighbors
   in
       if altitude lowest ≥ altitude (x, y)
             then Nothing
             else just lowest

 printArray arr =
     let ((sx, sy), (ex, ey)) = bounds arr
         printRow y = forM_ [sx..ex] (λ x → putStr (show (arr! (x, y)) ++ "\ t"))
     in
         forM_ [sy..ey] (λ y → printRow y "putStrLn" ")

 main = 
   do printArray altitudes
      putStrLn "- drains -------------"
      printArray drains
      putStrLn "- drains ′ ------------"
      printArray drains ′

 - simple recursion -

 drains = array cells [(cell, drainOf cell) |  cell ← range cells]

 drainOf cell |  trace ("drainOf" ++ show cell) False = undefined
 drainOf cell = 
   case lowestNeighbourOf cell of
     Nothing → cell
     Just nb → drainOf nb

 - lazy caching -

 drains ′ = array cells [(cell, drainOf ′ cell) |  cell ← range cells]

 drainOf ′ cell |  trace ("drainOf ′" ++ show cell) False = undefined
 drainOf ′ cell = 
   case lowestNeighbourOf cell of
     Nothing → cell
     Just nb → drains ′!  nb

Result:

 10 8 10	
 15 3 15	
 20 25 20	
 - drains -------------
 drainOf (1,1)
 drainOf (2.1)
 drainOf (2.2)
 drainOf (2.1)
 drainOf (2.2)
 drainOf (3.1)
 drainOf (2.1)
 drainOf (2.2)
 (2.2) (2.2) (2.2)	
 drainOf (1,2)
 drainOf (2.2)
 drainOf (2.2)
 drainOf (3.2)
 drainOf (2.2)
 (2.2) (2.2) (2.2)	
 drainOf (1,3)
 drainOf (1,2)
 drainOf (2.2)
 drainOf (2,3)
 drainOf (2.2)
 drainOf (3.3)
 drainOf (3.2)
 drainOf (2.2)
 (2.2) (2.2) (2.2)	
 - drains' ------------
 drainOf '(1,1)
 drainOf '(2,1)
 drainOf (2.2)
 drainOf '(3,1)
 (2.2) (2.2) (2.2)	
 drainOf (1,2)
 drainOf '(3,2)
 (2.2) (2.2) (2.2)	
 drainOf (1,3)
 drainOf (2,3)
 drainOf '(3.3)
 (2.2) (2.2) (2.2)	

As you can see, the first variant calculated the drain of each cell again and again, whereas the second variant did it exactly once.