fork download
  1. from collections import deque
  2. import numpy as np
  3.  
  4. asp = np.array([[ 2, 1, 8, 2],
  5. [ 5, 4, 6, 4],
  6. [ 3, 3, 3, 3],
  7. [ 2, 1, 8, 2]]) #aspect
  8. slp = np.array([[11,11,12,11],
  9. [12,11,12,10],
  10. [ 9, 9, 9,10],
  11. [11,10,11,11]]) #slope
  12. elv = np.array([[16,15,15,16],
  13. [17,15,15,15],
  14. [15,15,14,14],
  15. [16,15,16,16]]) #elevation
  16.  
  17. criteria = np.stack((asp, slp, elv))
  18.  
  19. def floodfill(n, start):
  20. def check(r, c):
  21. return result[r, c] == 0 and np.array_equal(crit, criteria[:, r, c])
  22.  
  23. stack = deque([start])
  24. crit = criteria [(slice(None), *start)]
  25. while stack:
  26. r, c = stack.popleft()
  27. if result[r, c]: continue
  28. result[r, c] = n
  29. if r > 0 and check(r - 1, c):
  30. stack.append((r - 1, c))
  31. if r < result.shape[0] - 1 and check(r + 1, c):
  32. stack.append((r + 1, c))
  33. if c > 0 and check(r, c - 1):
  34. stack.append((r, c - 1))
  35. if c < result.shape[1] - 1 and check(r, c + 1):
  36. stack.append((r, c + 1))
  37.  
  38.  
  39. result = np.zeros_like(asp)
  40. count = 1
  41. it = np.nditer(result, flags=['multi_index'])
  42. for x in it:
  43. if x == 0:
  44. floodfill(count, it.multi_index)
  45. count += 1
  46.  
  47. print(result)
Success #stdin #stdout 0.07s 92224KB
stdin
Standard input is empty
stdout
[[ 1  2  3  4]
 [ 5  6  7  8]
 [ 9  9 10 11]
 [12 13 14 15]]