O(klogk), priority queue, c++ solution


  • 1
    public:
        int kthSmallest(vector<vector<int>>& matrix, int k) {
            if(matrix.empty())
                return 0;
            int m=matrix.size();
            int n=matrix[0].size();
            auto cmp =[&matrix](pair<int,int> l,pair<int,int> r) {return matrix[l.first][l.second]>matrix[r.first][r.second];};
            priority_queue<pair<int,int>,vector<pair<int,int>>,decltype(cmp)> hp(cmp);
            vector<vector<int>> dict(m,vector<int> (n,1));
            pair<int,int> start={0,0};
            dict[0][0]=0;
            hp.push(start);
            while(!hp.empty()){
                pair<int,int> tmp=hp.top();
                hp.pop();
                k--;
                if(k==0)
                    return matrix[tmp.first][tmp.second];
                if(tmp.first+1<m&&dict[tmp.first+1][tmp.second]){
                    hp.push({tmp.first+1,tmp.second});
                    dict[tmp.first+1][tmp.second]=0;
                }
                if(tmp.second+1<n&&dict[tmp.first][tmp.second+1]){
                    hp.push({tmp.first,tmp.second+1});
                    dict[tmp.first][tmp.second+1]=0;
                }
            }
            return INT_MIN;
    
    }
    

    };


  • 0
    B

    Similar idea but in python

    def kthSmallest(self, matrix, k):
            n = len(matrix)
            hq = [(v,0,i) for i,v in enumerate(matrix[0])]
            heapq.heapify(hq)
            for _ in range(k-1):
                item = heapq.heappop(hq)
                if item[1] ==  n-1: continue
                else:
                    heapq.heappush(hq, (matrix[item[1]+1][item[2]], item[1]+1, item[2]))
            kthitem = heapq.heappop(hq)
            return kthitem[0]
    

  • 0
    J

    @blackbird, it looks like this is actually O(n + k log k) and not O(k log k) since you heapify n elements in hq.


  • 0
    B

    @jchennn Yes. Its actually O(n+klogn). What I mean by similar idea is that both solutions have min heap in it. Since k is probably greater than n, I think there's no need to implement a more complicated solution.


  • 1
    A

    Aren't you incurring O(n^2) time and space when you call

    vector<vector<int>> dict(m,vector<int> (n,1));
    

    There are other ways to mark matrix entries as "visited" that don't kill your time complexity.


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.