1706 Where Will the Ball Fall

이 문제도 조금 색다르긴 한데 어쨌든 전형적인 DFS 문제이다.

 

C++

class Solution {
public:
    vector<int> findBall(vector<vector<int>>& grid) {
        vector<int> res(grid[0].size());
        for(size_t i = 0; i < res.size(); ++i){
            size_t j = i, s = 0;
            while(res[i] == 0 && s < grid.size()){
                if(grid[s][j] == 1){
                    if (j == res.size() - 1 || grid[s][j + 1] == -1)
                        res[i] = -1;
                    else{
                        ++j;
                        ++s;
                    }
                }
                else if (grid[s][j] == -1){
                    if (j == 0 || grid[s][j - 1] == 1)
                        res[i] = -1;
                    else{
                        --j;
                        ++s;
                    }
                }
            }
            if (res[i] == 0)
                res[i] = j;
        }
        return res;
    }
};

Python

class Solution:
    def findBall(self, grid: List[List[int]]) -> List[int]:
        res = [0 for i in range(len(grid[0]))]
        
        for i in range(len(res)):
            j = i; s= 0;
            while res[i] == 0 and s < len(grid):
                if grid[s][j] == 1:
                    if j == len(res) - 1 or grid[s][j + 1] == -1:
                        res[i] = -1
                    else:
                        j += 1
                        s += 1
                elif grid[s][j] == -1:
                    if j == 0 or grid[s][j - 1] == 1:
                        res[i] = -1
                    else:
                        j -= 1
                        s += 1
            if res[i] == 0:
                res[i] = j
        
        return res