8ms Solution, find median with quick partition


  • 0
    C

    Similar to other find median base solutions, just changed to nth_element to do a quick partition to find median and use saved row and column indices to calculate distance without going over the entire grid again. The latter made performance better with large sparse grid, where #person/#gridPoint ratio is small.

    Time complexity still O(mn), space complexity, O(K), K being number of people in the grid, worse case space is O(mn).

    class Solution {
    public:
        int minTotalDistance(vector<vector<int>>& grid) {
            int n = grid.size();
            if (n == 0) return 0;
            int m = grid[0].size();
            if (m == 0) return 0;
            
            vector<int> row;
            vector<int> column;
            for (int i=0; i<n; i++) {
                for (int j=0; j<m; j++) {
                    if (grid[i][j] == 1) {
                        row.push_back(i);
                        column.push_back(j);
                    }
                }
            }
            
            nth_element(row.begin(), row.begin()+row.size()/2, row.end());
            nth_element(column.begin(), column.begin()+column.size()/2, column.end());
            
            int mi = row[row.size()/2];
            int mj = column[column.size()/2];
            
            long long distance = 0;
            for (auto i : row) distance += abs(i - mi);
            for (auto j : column) distance += abs(j - mj);
            
            return distance;
        }
    };
    

Log in to reply
 

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