Tìm trung vị trong ma trận đã sắp xếp theo hàng trong C++

Mô tả vấn đề

Cho ma trận 2D mat[r][c] có các phần tử được sắp xếp theo hàng (row-wise). Nhiệm vụ là tìm trung vị (median) của tất cả các phần tử trong ma trận.

Ví dụ minh họa

Đầu vào:

mat = { {2, 4, 7}, {5, 6, 8}, {4, 8, 9} }

Đầu ra:

6

Giải thích:

Các phần tử gộp lại: {2, 4, 4, 5, 6, 7, 8, 8, 9}
Trung vị (phần tử giữa) là 6.

Phương pháp giải quyết

Cách đơn giản: Gộp tất cả phần tử vào mảng 1D, sắp xếp và lấy phần tử giữa. Độ phức tạp O(rc log rc).

Cách hiệu quả (Tìm kiếm nhị phân): Trung vị có đúng (r*c)/2 phần tử nhỏ hơn nó. Ta dùng tìm kiếm nhị phân trên khoảng giá trị [min, max] của ma trận. Với mỗi giá trị mid, đếm số phần tử <= mid (dùng upper_bound trên mỗi hàng vì hàng đã sắp xếp). Điều chỉnh khoảng tìm kiếm dựa trên số lượng đếm được.

Chương trình minh họa

#include <bits/stdc++.h>
using namespace std;

#define r 3
#define c 3

int findMedian(int mat[][c]) {
    int smallest = INT_MAX, largest = INT_MIN;
    for (int i = 0; i < r; i++) {
        if (mat[i][0] < smallest) smallest = mat[i][0];
        if (mat[i][c-1] > largest) largest = mat[i][c-1];
    }
    while (smallest < largest) {
        int mid = smallest + (largest - smallest) / 2;
        int smallCount = 0;
        for (int i = 0; i < r; ++i)
            smallCount += upper_bound(mat[i], mat[i]+c, mid) - mat[i];
        if (smallCount < ((r * c + 1) / 2))
            smallest = mid + 1;
        else
            largest = mid;
    }
    return smallest;
}

int main() {
    int mat[][c] = { {2, 5, 7}, {4, 6, 8}, {1, 8, 9} };
    cout << "Trung vị của ma trận là " << findMedian(mat);
    return 0;
}

Kết quả đầu ra

Trung vị của ma trận là 6