ソートに関する練習問題！

問題へのリンク

問題概要

長さ $N$ の数列 $A_{1}, A_{2}, \dots, A_{N}$ と、長さ $M$ の数列 $B_{1}, B_{2}, \dots, B_{M}$ が与えられる。これらの数列の値はすべて互いに相異なる。

これらの 2 つの数列の各要素 $v$ について、次の値を求めよ。

2 つの数列を連結してできる長さ $N+M$ の数列を小さい順に並べ替える。

このとき、値 $v$ がその並び替えた数列において何番目にあるか。

制約

$1 \le N, M \le 10^{5}$

考えたこと

いきなり 2 つの数列を考えるのはややこしいので、1 つの数列で考えよう。つまり、数列 $A_{1}, A_{2}, \dots, A_{N}$ の各要素について、その要素が小さい順に何番目の値であるかを求めよう。

これをするためには、次のペア値を要素にもつ配列を考えると良い。

(値, 数列 $A$ における添字)

たとえば、 $A = (6, 1, 9, 2, 11)$ のときは、次の配列を考える。

$\lbrack (6, 1), (1, 2), (9, 3), (2, 4), (11, 5) \rbrack$

これを小さい順に並び替えよう。そうすると、次のようになる。

$\lbrack (1, 2), (2, 4), (6, 1), (9, 3), (11, 5) \rbrack$

これを元に、数列 $A$ の各要素が小さい順に何番目であるかが求められる。たとえば、並び替えたあとの 4 番目の要素が (9, 3) であるば、これは

「 $A_{3}$ が小さい順に 4 番目であること」

を表している。あとは、以上の処理をプログラムに起こせばよい。このソート後のペア値の配列を $C$ とすると、次のように書ける (0-indexed で書いている)。

vector<int> res(N);
for (int i = 0; i < N; ++i) {
    // 数列 A の C[i].second 番目の要素は、小さい順に i 番目である
    res[C[i].second] = i;
}

ここまでは数列が 1 個の場合を考えてきたが、数列が 2 個になっても同様にできる。

計算量は $O( (N+M)\log(N+M) )$ となる。

コード

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

int main() {
    int N, M;
    cin >> N >> M;
    vector<int> A(N), B(M);
    for (int i = 0; i < N; ++i) cin >> A[i];
    for (int i = 0; i < M; ++i) cin >> B[i];
    
    // C を index つきで表す (B の index は N, N+1, ..., N+M-1 とする)
    using pint = pair<int, int>;  // (値, index)
    vector<pint> C(N + M);
    for (int i = 0; i < N; ++i) C[i] = {A[i], i};
    for (int i = 0; i < M; ++i) C[i + N] = {B[i], i + N};
    
    // C をソート
    sort(C.begin(), C.end());
    
    // C の値が小さい順に見ていく
    vector<int> res(N + M);
    for (int i = 0; i < N+M; ++i) {
        // 添字が C[i].second である要素が、i 番目に小さい要素である
        res[C[i].second] = i;
    }
    
    // 出力 (1-indexed に直す)
    for (int i = 0; i < N; ++i) cout << res[i]+1 << " ";
    cout << endl;
    for (int i = 0; i < M; ++i) cout << res[i+N]+1 << " ";
    cout << endl;
}

けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 294 C - Merge Sequences (5Q, 灰色, 300 点)

問題概要

制約

考えたこと

コード