結果
| 問題 |
No.3078 Difference Sum Query
|
| ユーザー |
 atcoder8
|
| 提出日時 | 2025-03-31 07:22:05 |
| 言語 | Rust (1.83.0 + proconio) |
| 結果 |
AC
|
| 実行時間 | 79 ms / 2,000 ms |
| コード長 | 8,592 bytes |
| コンパイル時間 | 14,257 ms |
| コンパイル使用メモリ | 398,916 KB |
| 実行使用メモリ | 22,400 KB |
| 最終ジャッジ日時 | 2025-03-31 07:22:27 |
| 合計ジャッジ時間 | 19,981 ms |
|
ジャッジサーバーID （参考情報） |
judge2 / judge1 |
（要ログイン）
| ファイルパターン | 結果 |
|---|---|
| sample | AC * 2 |
| other | AC * 26 |
ソースコード
use std::ops::{AddAssign, Range, Sub};
use fenwick_tree::FenwickTree;
use proconio::{input, marker::Usize1};
fn main() {
input! {
(n, q): (usize, usize),
aa: [i64; n],
lrx: [(Usize1, usize, i64); q],
}
let mut ia = aa.iter().cloned().enumerate().collect::<Vec<_>>();
ia.sort_unstable_by_key(|&(_, a)| a);
let mut queries = lrx
.iter()
.enumerate()
.map(|(i, &(l, r, x))| Query {
id: i,
range: l..r,
x,
})
.collect::<Vec<_>>();
queries.sort_unstable_by_key(|query| query.x);
let mut ft = FenwickTree::<Node>::new(n);
for (i, &a) in aa.iter().enumerate() {
ft.add(
i,
Node {
sum_a: a,
num_positive_x: 0,
},
);
}
let mut answers = vec![0_i64; q];
let mut progress = 0_usize;
for query in &queries {
while progress < n && ia[progress].1 < query.x {
let (i, a) = ia[progress];
ft.add(
i,
Node {
sum_a: -2 * a,
num_positive_x: 1,
},
);
progress += 1;
}
let Node {
sum_a,
num_positive_x,
} = ft.sum(query.range.clone());
answers[query.id] = sum_a + query.x * (2 * num_positive_x - query.range.len() as i64);
}
println!(
"{}",
answers
.iter()
.map(ToString::to_string)
.collect::<Vec<_>>()
.join("\n")
);
}
#[derive(Debug, Clone)]
struct Query {
id: usize,
range: Range<usize>,
x: i64,
}
#[derive(Debug, Clone, Copy, Default)]
struct Node {
sum_a: i64,
num_positive_x: i64,
}
impl AddAssign for Node {
fn add_assign(&mut self, rhs: Self) {
self.sum_a += rhs.sum_a;
self.num_positive_x += rhs.num_positive_x;
}
}
impl Sub for Node {
type Output = Self;
fn sub(self, rhs: Self) -> Self::Output {
Self::Output {
sum_a: self.sum_a - rhs.sum_a,
num_positive_x: self.num_positive_x - rhs.num_positive_x,
}
}
}
pub mod fenwick_tree {
//! Processes the following query in `O(log(n))` time
//! for a sequence of numbers with `n` elements:
//! * Update one element
//! * Calculate the sum of the elements of a range
//! * Gets the elements of a number sequence.
use std::ops::{AddAssign, RangeBounds, Sub, SubAssign};
/// # Examples
///
/// ```
/// use atcoder8_library::fenwick_tree::FenwickTree;
///
/// let ft = FenwickTree::from(vec![3, -1, 4, 1, -5, 9, 2]);
/// assert_eq!(ft.sum(2..), 11);
/// ```
#[derive(Debug, Clone)]
pub struct FenwickTree<T>(Vec<T>);
impl<T> From<Vec<T>> for FenwickTree<T>
where
T: Default + Clone + AddAssign<T>,
{
/// # Examples
///
/// ```
/// use atcoder8_library::fenwick_tree::FenwickTree;
///
/// let ft = FenwickTree::from(vec![3, -1, 4, 1, -5, 9, 2]);
/// assert_eq!(ft.sum(2..6), 9);
/// ```
fn from(t: Vec<T>) -> Self {
let mut ft = FenwickTree::new(t.len());
for (i, x) in t.into_iter().enumerate() {
ft.add(i, x);
}
ft
}
}
impl<T> FenwickTree<T> {
/// Constructs a `FenwickTree<T>` with `n` elements.
///
/// Each element is initialized with `T::default()`.
///
/// # Examples
///
/// ```
/// use atcoder8_library::fenwick_tree::FenwickTree;
///
/// let ft = FenwickTree::<i32>::new(5);
/// assert_eq!(ft.sum(..), 0);
/// ```
pub fn new(n: usize) -> Self
where
T: Default + Clone,
{
FenwickTree(vec![T::default(); n])
}
/// Add `x` to the `p`-th element.
///
/// # Examples
///
/// ```
/// use atcoder8_library::fenwick_tree::FenwickTree;
///
/// let mut ft = FenwickTree::from(vec![3, -1, 4, 1, -5, 9, 2]);
/// assert_eq!(ft.sum(2..6), 9);
///
/// ft.add(3, 100);
/// assert_eq!(ft.sum(2..6), 109);
/// ```
pub fn add(&mut self, p: usize, x: T)
where
T: Clone + AddAssign<T>,
{
let FenwickTree(data) = self;
let n = data.len();
assert!(p < n);
let mut p = p + 1;
while p <= n {
data[p - 1] += x.clone();
p += p & p.overflowing_neg().0;
}
}
/// Subtract `x` from the `p`-th element.
///
/// # Examples
///
/// ```
/// use atcoder8_library::fenwick_tree::FenwickTree;
///
/// let mut ft = FenwickTree::<u32>::from(vec![3, 1, 4, 1, 5, 9, 2]);
/// assert_eq!(ft.sum(2..6), 19);
///
/// ft.sub(3, 1);
/// assert_eq!(ft.sum(2..6), 18);
/// ```
pub fn sub(&mut self, p: usize, x: T)
where
T: Clone + SubAssign<T>,
{
let FenwickTree(data) = self;
let n = data.len();
assert!(p < n);
let mut p = p + 1;
while p <= n {
data[p - 1] -= x.clone();
p += p & p.overflowing_neg().0;
}
}
/// Sets `x` to the `p`-th element.
///
/// # Examples
///
/// ```
/// use atcoder8_library::fenwick_tree::FenwickTree;
///
/// let mut ft = FenwickTree::from(vec![3, -1, 4, 1, -5, 9, 2]);
/// assert_eq!(ft.sum(2..6), 9);
///
/// ft.set(3, 100);
/// assert_eq!(ft.sum(2..6), 108);
/// ```
pub fn set(&mut self, p: usize, x: T)
where
T: Default + Clone + AddAssign<T> + Sub<T, Output = T> + SubAssign<T>,
{
let FenwickTree(data) = self;
let n = data.len();
assert!(p < n);
self.sub(p, self.get(p));
self.add(p, x);
}
/// Compute the sum of the range [0, r).
fn inner_sum(&self, r: usize) -> T
where
T: Default + Clone + AddAssign<T>,
{
let mut s = T::default();
let mut r = r;
while r > 0 {
s += self.0[r - 1].clone();
r -= r & r.wrapping_neg();
}
s
}
/// Calculate the total of the range.
///
/// # Examples
///
/// ```
/// use atcoder8_library::fenwick_tree::FenwickTree;
///
/// let ft = FenwickTree::from(vec![3, -1, 4, 1, -5, 9, 2]);
/// assert_eq!(ft.sum(..), 13);
/// assert_eq!(ft.sum(2..), 11);
/// assert_eq!(ft.sum(..6), 11);
/// assert_eq!(ft.sum(2..6), 9);
/// assert_eq!(ft.sum(6..2), 0);
/// ```
pub fn sum<R>(&self, rng: R) -> T
where
T: Default + Clone + AddAssign<T> + Sub<T, Output = T>,
R: RangeBounds<usize>,
{
let n = self.0.len();
let l = match rng.start_bound() {
std::ops::Bound::Included(&start_bound) => start_bound,
std::ops::Bound::Excluded(&start_bound) => start_bound + 1,
std::ops::Bound::Unbounded => 0,
};
let r = match rng.end_bound() {
std::ops::Bound::Included(&end_bound) => end_bound + 1,
std::ops::Bound::Excluded(&end_bound) => end_bound,
std::ops::Bound::Unbounded => n,
};
assert!(l <= n && r <= n);
if l >= r {
T::default()
} else {
self.inner_sum(r) - self.inner_sum(l)
}
}
/// Returns the value of an element in a sequence of numbers.
/// Calculate the total of the range.
///
/// # Examples
///
/// ```
/// use atcoder8_library::fenwick_tree::FenwickTree;
///
/// let ft = FenwickTree::from(vec![3, -1, 4, -1, 5]);
/// assert_eq!(ft.get(2), 4);
/// ```
pub fn get(&self, p: usize) -> T
where
T: Default + Clone + AddAssign<T> + Sub<T, Output = T>,
{
assert!(p < self.0.len());
self.sum(p..=p)
}
}
}