# 1706. Where Will the Ball Fall

• Time: $O(mn)$
• Space: $O(n)$
  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 class Solution { public: vector findBall(vector>& grid) { const int m = grid.size(); const int n = grid[0].size(); // dp[i] := status of i-th column // -1 := empty, 0 := b0, 1 := b1, ... vector dp(n); // ans[i] := i-th ball's final position vector ans(n, -1); iota(begin(dp), end(dp), 0); for (int i = 0; i < m; ++i) { vector newDp(n, -1); for (int j = 0; j < n; ++j) { // out of bound if (j + grid[i][j] < 0 || j + grid[i][j] == n) continue; // stuck if (grid[i][j] == 1 && grid[i][j + 1] == -1 || grid[i][j] == -1 && grid[i][j - 1] == 1) continue; newDp[j + grid[i][j]] = dp[j]; } dp = move(newDp); } for (int i = 0; i < n; ++i) if (dp[i] != -1) ans[dp[i]] = i; return ans; } }; 
  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 class Solution { public int[] findBall(int[][] grid) { final int m = grid.length; final int n = grid[0].length; // dp[i] := status of i-th column // -1 := empty, 0 := b0, 1 := b1, ... int[] dp = new int[n]; // ans[i] := i-th ball's final position int[] ans = new int[n]; Arrays.fill(ans, -1); for (int i = 0; i < n; ++i) dp[i] = i; for (int i = 0; i < m; ++i) { int[] newDp = new int[n]; Arrays.fill(newDp, -1); for (int j = 0; j < n; ++j) { // out of bound if (j + grid[i][j] < 0 || j + grid[i][j] == n) continue; // stuck if (grid[i][j] == 1 && grid[i][j + 1] == -1 || grid[i][j] == -1 && grid[i][j - 1] == 1) continue; newDp[j + grid[i][j]] = dp[j]; } dp = newDp; } for (int i = 0; i < n; ++i) if (dp[i] != -1) ans[dp[i]] = i; return ans; } } 
  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 class Solution: def findBall(self, grid: List[List[int]]) -> List[int]: m = len(grid) n = len(grid[0]) # dp[i] := status of i-th column # -1 := empty, 0 := b0, 1 := b1, ... dp = [i for i in range(n)] # ans[i] := i-th ball's final positio ans = [-1] * n for i in range(m): newDp = [-1] * n for j in range(n): # out of bound if j + grid[i][j] < 0 or j + grid[i][j] == n: continue if grid[i][j] == 1 and grid[i][j + 1] == -1 or \ grid[i][j] == -1 and grid[i][j - 1] == 1: continue newDp[j + grid[i][j]] = dp[j] dp = newDp for i, ball in enumerate(dp): if ball != -1: ans[ball] = i return ans