Skip to content

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<int> findBall(vector<vector<int>>& grid) {
    const int m = grid.size();
    const int n = grid[0].size();
    // dp[i] := status of the i-th column
    // -1 := empty, 0 := b0, 1 := b1, ...
    vector<int> dp(n);
    // ans[i] := the i-th ball's final position
    vector<int> ans(n, -1);

    iota(dp.begin(), dp.end(), 0);

    for (int i = 0; i < m; ++i) {
      vector<int> newDp(n, -1);
      for (int j = 0; j < n; ++j) {
        // out-of-bounds
        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 = std::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 the i-th column
    // -1 := empty, 0 := b0, 1 := b1, ...
    int[] dp = new int[n];
    // ans[i] := the 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-bounds
        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 the i-th column
    # -1 := empty, 0 := b0, 1 := b1, ...
    dp = [i for i in range(n)]
    # ans[i] := the i-th ball's final positio
    ans = [-1] * n

    for i in range(m):
      newDp = [-1] * n
      for j in range(n):
        # out-of-bounds
        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