Two Dimensional Waterflow problem

Question

The one dimensional twitter waterflow problem is this:

You are given an array that represents a hill in the sense that the ith entry is the height of the ith location of the hill. When it rains, water logs in the hills, and you need to figure out how much water would log.

For example, after raining, the array 2 5 3 4 3 2 5 5 3 4 2 2 2 looks like this,

and 9 units of water have accumulated.

---

The challenge is to solve the 2D version of this. To clarify, you will be given a 2D matrix of land heights, and you will have to output the total amount of rainwater it can collect. Water cannot flow out through diagonals, but only the four cardinal directions.

Answer 1

Python 3, O(h × w × d) 296 bytes, O(h × w × d')

def f(a):
 w=0;b=[len(a[0])*[0]];a=b+a+b;a=[[0]+x+[0]for x in a];b=sorted(set(sum(a,[])))
 for[i,j]in zip(b,b[1:]+[1+max(b)]):
  c=[[0,0]]
  while c:
   [x,y]=c.pop()
   if x<len(a)and x>=0and y<len(a[0])and y>=0and a[x][y]<=i:w+=i-a[x][y];a[x][y]=j;c+=[[x+1,y],[x,y+1],[x-1,y],[x,y-1]]
 return w

Try it online! Where d' represents the number of distinct values in the input. Edit: Saved 3 bytes thanks to @Mr.Xcoder.