AtCoder Beginner Contest 412 F

https://atcoder.jp/contests/abc412/tasks/abc412_f

外にある靴下も合わせてAにして降順にソートします。今外にある靴下がでを取り出したとき、なら終わり、ならをそのままに、ならをそのままにすればよいので、期待値は、

として、

となります。これは少し工夫すれば解けます。
ただ、剰余を計算するコードを書くのが気を遣わないといけないくていつも辛いので、ふつうに計算できる構造体を作ります。ただし、倍くらい時間がかかっています。

// Socks 4
#![allow(non_snake_case)]


//////////////////// library ////////////////////

fn read<T: std::str::FromStr>() -> T {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).ok();
    line.trim().parse().ok().unwrap()
}

fn read_vec<T: std::str::FromStr>() -> Vec<T> {
    read::<String>().split_whitespace()
            .map(|e| e.parse().ok().unwrap()).collect()
}

// ax = by + 1 (a, b > 0)
fn linear_diophantine(a: i64, b: i64) -> Option<(i64, i64)> {
    if a == 1 {
        return Some((1, 0))
    }
    
    let q = b / a;
    let r = b % a;
    if r == 0 {
        return None
    }
    let (x1, y1) = linear_diophantine(r, a)?;
    Some((-q * x1 - y1, -x1))
}

fn inverse<const D: i64>(a: i64) -> i64 {
    let (x, _y) = linear_diophantine(a, D).unwrap();
    x.rem_euclid(D)
}


//////////////////// ModInt ////////////////////

use std::cmp::Ordering;
use std::ops::{Add, AddAssign, Sub, SubAssign,
               Mul, MulAssign, Div, DivAssign, Neg};

#[derive(Copy, Clone)]
struct ModInt<const D: i64>(i64);

impl<const D: i64> From<i64> for ModInt<D> {
    fn from(x: i64) -> Self {
        Self(x.rem_euclid(D))
    }
}
impl<const D: i64> PartialEq for ModInt<D> {
    fn eq(&self, other: &Self) -> bool {
        self.0 == other.0
    }
}

impl<const D: i64> Eq for ModInt<D> {}

impl<const D: i64> PartialOrd for ModInt<D> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl<const D: i64> Ord for ModInt<D> {
    fn cmp(&self, other: &Self) -> Ordering {
        self.0.cmp(&other.0)
    }
}

impl<const D: i64> Add for ModInt<D> {
    type Output = Self;
    
    fn add(self, other: Self) -> Self {
        Self((self.0 + other.0) % D)
    }
}

impl<const D: i64> Add<i64> for ModInt<D> {
    type Output = Self;
    
    fn add(self, other: i64) -> Self {
        Self((self.0 + other) % D)
    }
}

impl<const D: i64> Add<ModInt<D>> for i64 {
    type Output = ModInt<D>;
    
    fn add(self, other: ModInt<D>) -> ModInt<D> {
        ModInt::<D>((self + other.0).rem_euclid(D))
    }
}

impl<const D: i64> AddAssign for ModInt<D> {
    fn add_assign(&mut self, other: ModInt<D>) {
        self.0 = (self.0 + other.0).rem_euclid(D)
    }
}

impl<const D: i64> AddAssign<i64> for ModInt<D> {
    fn add_assign(&mut self, other: i64) {
        self.0 = (self.0 + other).rem_euclid(D)
    }
}

impl<const D: i64> Sub for ModInt<D> {
    type Output = Self;
    
    fn sub(self, other: Self) -> Self {
        Self((self.0 - other.0).rem_euclid(D))
    }
}

impl<const D: i64> Sub<i64> for ModInt<D> {
    type Output = Self;
    
    fn sub(self, other: i64) -> Self {
        Self((self.0 - other).rem_euclid(D))
    }
}

impl<const D: i64> Sub<ModInt<D>> for i64 {
    type Output = ModInt<D>;
    
    fn sub(self, other: ModInt<D>) -> ModInt<D> {
        ModInt::<D>((self - other.0).rem_euclid(D))
    }
}

impl<const D: i64> SubAssign for ModInt<D> {
    fn sub_assign(&mut self, other: Self) {
        self.0 = (self.0 - other.0).rem_euclid(D)
    }
}

impl<const D: i64> SubAssign<i64> for ModInt<D> {
    fn sub_assign(&mut self, other: i64) {
        self.0 = (self.0 - other).rem_euclid(D)
    }
}

impl<const D: i64> Mul for ModInt<D> {
    type Output = Self;
    
    fn mul(self, other: Self) -> Self {
        Self(self.0 * other.0 % D)
    }
}

impl<const D: i64> Mul<i64> for ModInt<D> {
    type Output = Self;
    
    fn mul(self, other: i64) -> Self {
        Self(self.0 * other % D)
    }
}

impl<const D: i64> Mul<ModInt<D>> for i64 {
    type Output = ModInt<D>;
    
    fn mul(self, other: ModInt<D>) -> ModInt<D> {
        ModInt::<D>(self * other.0 % D)
    }
}

impl<const D: i64> MulAssign for ModInt<D> {
    fn mul_assign(&mut self, other: Self) {
        self.0 = self.0 * other.0 % D
    }
}

impl<const D: i64> MulAssign<i64> for ModInt<D> {
    fn mul_assign(&mut self, other: i64) {
        self.0 = self.0 * other % D
    }
}

impl<const D: i64> Div for ModInt<D> {
    type Output = Self;
    
    fn div(self, other: Self) -> Self {
        Self(self.0 * inverse::<D>(other.0) % D)
    }
}

impl<const D: i64> Div<i64> for ModInt<D> {
    type Output = Self;
    
    fn div(self, other: i64) -> Self {
        Self(self.0 * inverse::<D>(other) % D)
    }
}

impl<const D: i64> Div<ModInt<D>> for i64 {
    type Output = ModInt<D>;
    
    fn div(self, other: ModInt<D>) -> ModInt<D> {
        ModInt::<D>(self * inverse::<D>(other.0) % D)
    }
}

impl<const D: i64> DivAssign for ModInt<D> {
    fn div_assign(&mut self, other: Self) {
        self.0 = self.0 * inverse::<D>(other.0) % D
    }
}

impl<const D: i64> DivAssign<i64> for ModInt<D> {
    fn div_assign(&mut self, other: i64) {
        self.0 = self.0 * inverse::<D>(other) % D
    }
}

impl<const D: i64> Neg for ModInt<D> {
    type Output = Self;
    
    fn neg(self) -> Self::Output {
        Self(if self.0 == 0 { 0 } else { D - self.0 })
    }
}

use std::iter::Sum;

impl<const D: i64> Sum for ModInt<D> {
    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
        iter.fold(Self::from(0), |acc, x| acc + x)
    }
}


//////////////////// process ////////////////////

const D: i64 = 998244353;

type IntD = ModInt::<D>;

fn read_input() -> (usize, Vec<IntD>) {
    let v: Vec<usize> = read_vec();
    let C = v[1] - 1;
    let A: Vec<i64> = read_vec();
    let B: Vec<IntD> = A.into_iter().map(IntD::from).collect();
    (C, B)
}

fn F(C: usize, mut A: Vec<IntD>) -> i64 {
    A[C] += 1;
    let mut v: Vec<(IntD, usize)> = A.iter().cloned().zip(0..).collect();
    v.sort_by(|a, b| b.cmp(&a));
    let first = v.iter().enumerate().filter(|&(_, t)| t.1 == C).
                                        map(|(i, _)| i).next().unwrap();
    A.sort_by(|a, b| b.cmp(&a));
    let S = A.iter().cloned().sum::<IntD>() - 1;
    let mut B = IntD::from(0);
    let mut C = 1 - (A[0] - 1) / S;
    for i in 0.. {
        let E = (1 + B) / (1 - C);
        if i == first {
            return E.0
        }
        B += A[i] * E / S;
        C -= A[i+1] / S
    }
    0
}

fn main() {
    let (C, A) = read_input();
    println!("{}", F(C, A))
}