Trick 5: how to get rid of TLE in BFS

5314. Jump Game IV

Naive BFS will get TLE, traverse nodes reversely can get AC as our target is last element of the array and do the searching reversely will make the similar effect of A* algorithm. We search from the close to target element firstly.

Reversely works very good on the test case below

[7,…,7,7,11] it has many 7s at the beginning of the list.

Python code is here:

The critical part to get rid of TLE is using

a_i[arr[i]][::-1]

instead of a_i[arr[i]].

class Solution:
    def minJumps(self, arr: List[int]) -> int:
        a_i = collections.defaultdict(list)
        for i, a in enumerate(arr):
            a_i[a].append(i)
        def bfs():
            from collections import deque
            Q = deque([(0,0)])
            seen = {0}
            while Q:
                i, d = Q.popleft()
                if i==len(arr)-1:return d
                for j in [i-1, i+1]+a_i[arr[i]][::-1]:
                    if 0<=j<len(arr) and j!=i and j not in seen:
                        seen.add(j)
                        if j==len(arr)-1:return d+1
                        Q.append((j,d+1))
        return bfs()

C++ code is here.

class Solution {
     vector<int> dir = {-1, 1};
public:
     int bfs(int n,unordered_map<int, vector<int>> &a_i,vector<int>& arr){
            unordered_set<int> visited;
            deque<pair<int, int>> Q = {{0,0}};
            while(!Q.empty()){
                auto this_p = Q.front();
                Q.pop_front();
                int i=this_p.first, d= this_p.second;
                if (i == n-1)return d;
                for (auto &di:dir){
                    int j=i+di;
                    if(j>=0 && j<n){
                        //if(j==this_p.first)continue;
                        auto it=visited.find(j);
                        if(it==visited.end()){
                            if(j==n-1)return d+1;
                            visited.insert(j);
                            Q.emplace_back(j, d+1);
                    }
                }
                }
                
                for (auto &j:a_i[arr[i]]){
                    if(j==i)continue;
                    auto it=visited.find(j);
                    if(it==visited.end()){
                        if(j==n-1)return d+1;
                        visited.insert(j);
                        Q.emplace_back(make_pair(j, d+1));
                    }
                }
                
            }
         return 0;
        }
    int minJumps(vector<int>& arr) {
        unordered_map<int, vector<int>> a_i;
        int n = arr.size();
        /*
         // traverse in order get TLE at the fllowing case:
         [7,..7,11]
        for(int i=0;i<n;++i){
            a_i[arr[i]].push_back(i);
        }
        */
        
        // traverse reversely
        for(int i=n-1;i>=0;--i){
            a_i[arr[i]].push_back(i);
        }
        
       
        return bfs(n, a_i,arr);
    }
};

We can also get rid of TLE by early pruning.

Such as in the following example, the bolded 23 is unnecessary to be explored, we can safely remove it from candidates during early stage.

[100,-23,-23,404,100,23,23,23,3,404]

Python code

class Solution:
    def minJumps(self, arr: List[int]) -> int:
        a_i_,a_i = collections.defaultdict(list),collections.defaultdict(list)
        for i, a in enumerate(arr):
            a_i_[a].append(i)
        #early pruning.
        for key in a_i_:
            for j, v in enumerate(a_i_[key]):
                #print(j,v)
                if j>=1 and j<len(a_i_[key])-1 and a_i_[key][j-1]==v-1 and a_i_[key][j+1]==v+1:continue
                a_i[key].append(v)
        def bfs():
            from collections import deque
            Q = deque([(0,0)])
            seen = {0}
            while Q:
                i, d = Q.popleft()
                if i==len(arr)-1:return d
                for j in [i-1, i+1]+a_i[arr[i]]:
                    if 0<=j<len(arr) and j!=i and j not in seen:
                        seen.add(j)
                        if j==len(arr)-1:return d+1
                        Q.append((j,d+1))
        return bfs()

C++ code

class Solution {
     vector<int> dir = {-1, 1};
public:
     int bfs(int n,unordered_map<int, vector<int>> &a_i,vector<int>& arr){
            unordered_set<int> visited;
            deque<pair<int, int>> Q = {{0,0}};
            while(!Q.empty()){
                auto this_p = Q.front();
                Q.pop_front();
                int i=this_p.first, d= this_p.second;
                if (i == n-1)return d;
                for (auto &di:dir){
                    int j=i+di;
                    if(j>=0 && j<n){
                        //if(j==this_p.first)continue;
                        auto it=visited.find(j);
                        if(it==visited.end()){
                            if(j==n-1)return d+1;
                            visited.insert(j);
                            Q.emplace_back(j, d+1);
                    }
                }
                }
                
                for (auto &j:a_i[arr[i]]){
                    if(j==i)continue;
                    auto it=visited.find(j);
                    if(it==visited.end()){
                        if(j==n-1)return d+1;
                        visited.insert(j);
                        Q.emplace_back(make_pair(j, d+1));
                    }
                }
                
            }
         return 0;
        }
    int minJumps(vector<int>& arr) {
        unordered_map<int, vector<int>> a_i;
        int n = arr.size();
        /*
         // traverse in order get TLE at the fllowing case:
         [7,..7,11]
        for(int i=0;i<n;++i){
            a_i[arr[i]].push_back(i);
        }
        */
         
         // traverse in order get TLE, prune unnesessary nodes can get AC
         //[7,..7,11]
        unordered_map<int, vector<int>> a_i_;
        for(int i=0;i<n;++i){
            a_i_[arr[i]].push_back(i);
        }
        
        for(auto it=a_i_.begin();it!=a_i_.end();++it){
            int k=it->first;
            for(int j=0;j<a_i_[k].size();++j){
                if(j>=1 && j<a_i_[k].size()-1 && a_i_[k][j+1]==a_i_[k][j]+1 &&a_i_[k][j-1]==a_i_[k][j]-1)continue;
                a_i[k].push_back(a_i_[k][j]);
            }
        }
        /*
        // traverse reversely
        for(int i=n-1;i>=0;--i){
            a_i_[arr[i]].push_back(i);
        }
        */
        
       
        return bfs(n, a_i,arr);
    }
};