B - Better Students Are Needed! Editorial by cunitac
\(((1, A_1, B_1), (2, A_2, B_2), \ldots, (N, A_N, B_N))\) のような列を作り、
- 全体を \((-A_i, i)\) の昇順にソート
- \(X+1\) 番目以降を \((-B_i, i)\) の昇順にソート
- \(X+Y+1\) 番目以降を \((-(A_i+B_i), i)\) の昇順にソート
- \(X+Y+Z\) 番目以前を \(i\) の昇順にソート
した上で、\(X+Y+Z\) 番目以前を順に出力すればよいです。
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);
}
}
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])
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: