B - Better Students Are Needed! Editorial by cunitac


\(((1, A_1, B_1), (2, A_2, B_2), \ldots, (N, A_N, B_N))\) のような列を作り、

  1. 全体を \((-A_i, i)\) の昇順にソート
  2. \(X+1\) 番目以降を \((-B_i, i)\) の昇順にソート
  3. \(X+Y+1\) 番目以降を \((-(A_i+B_i), i)\) の昇順にソート
  4. \(X+Y+Z\) 番目以前を \(i\) の昇順にソート

した上で、\(X+Y+Z\) 番目以前を順に出力すればよいです。

実装例(Rust)

use itertools::izip;
use proconio::input;
use std::cmp::Reverse;

fn main() {
    input!(
        n: usize,
        x: usize,
        y: usize,
        z: usize,
        a: [u32; n],
        b: [u32; n]
    );

    let mut s = izip!(1.., a, b).collect::<Vec<_>>();
    s.sort_by_key(|&(i, a, _b)| (Reverse(a), i));
    s[x..].sort_by_key(|&(i, _a, b)| (Reverse(b), i));
    s[x + y..].sort_by_key(|&(i, a, b)| (Reverse(a + b), i));
    s[..x + y + z].sort_by_key(|&(i, ..)| i);

    for &(i, ..) in &s[..x + y + z] {
        println!("{}", i);
    }
}

実装例(Python)

n, x, y, z = map(int, input().split())
a = map(int, input().split())
b = map(int, input().split())

s = list(zip(range(1, n + 1), a, b))
s.sort(key=lambda p: (-p[1], p[0]))
s[x:] = sorted(s[x:], key=lambda p: (-p[2], p[0]))
s[x + y:] = sorted(s[x + y:], key=lambda p: (-(p[1] + p[2]), p[0]))
s[:x + y + z] = sorted(s[:x + y + z], key=lambda p: p[0])

for p in s[:x + y + z]:
  print(p[0])

実装例(C++)

tuple を扱うのは面倒なので、添字の配列をソートすることにします。また、pair をキーにソートするのは面倒なので、安定ソートをうまく用いることで回避します。たとえば、\((-A_i, i)\) の昇順にソートすることは、\(i\) の昇順にソートしてから \(A_i\) の降順に安定ソートすることと等価です。

#include <algorithm>
#include <iostream>
#include <numeric>
#include <vector>

using std::cin, std::cout, std::vector, std::sort, std::stable_sort;

int main() {
    int n, x, y, z;
    cin >> n >> x >> y >> z;
    vector<int> a(n), b(n);
    for (auto& e : a) cin >> e;
    for (auto& e : b) cin >> e;

    vector<int> s(n);
    std::iota(s.begin(), s.end(), 0);

    // 既に昇順なので不要
    // sort(s.begin(), s.end());
    stable_sort(s.begin(), s.end(), [&](int i, int j) { return a[i] > a[j]; });
    sort(s.begin() + x, s.end());
    stable_sort(s.begin() + x, s.end(), [&](int i, int j) { return b[i] > b[j]; });
    sort(s.begin() + x + y, s.end());
    stable_sort(s.begin() + x + y, s.end(), [&](int i, int j) { return a[i] + b[i] > a[j] + b[j]; });
    sort(s.begin(), s.begin() + x + y + z);

    for (int i = 0; i < x + y + z; i++)
        cout << (s[i] + 1) << '\n';
}

posted:
last update: