Matrix POJ3685

投稿日：2022/09/13 更新日:2022/09/13

問題

$N\times N$ の行列 $A$ があります。 $i$ 行目 $j$ 列目を $A_{ij}$ としたとき、その値は $i^2+100000\times i + j^2 - 100000\times j + i\times j$ に等しいです。整数 $M$ が与えられるので、 $M$ 番目に小さい値を出力してください。

解法

$A_{ij} \leqq v$ となる値が $M$ 個以上か、で二分探索します。

$A_{ij}$ の値は $j$ を固定した際、 $i$ について単調増加です。よって、各 $j$ に対し二分探索で $i$ の個数を数えて行けば良いです。

実装

#include <iostream>
using namespace std;

long long calc(long long i, long long j){
    return i*i + 100000*i + j*j - 100000*j + i*j;
}

void solve(){
    int n;
    long long m;
    cin >> n >> m;
    long long ok = 1e18, ng = -1e18;
    while(ok-ng > 1){
        long long v = (ok+ng)/2;

        long long count = 0;
        for(int j = 1;j <= n; ++j){
            int ok = 0, ng = n+1;
            while(ng - ok > 1){
                int mid = (ok+ng)/2;
                if(calc(mid, j) <= v){
                    ok = mid;
                }else{
                    ng = mid;
                }
            }
            count += ok;
        }

        if(count >= m){
            ok = v;
        }else{
            ng = v;
        }
    }
    cout << ok << endl;
}

int main(){
    int t;
    cin >> t;
    while(t--){
        solve();
    }
}