AtCoder 342

Result

One more time I managed to solve only A,B,C problems. I guess I am not very far away from solving the problem D consistently - I just need to get better at figuring out edge cases and algorithms details (e.g. “Should I add n+1 here or just n”?)

Problem B - Which is Ahead?

We have N persons standing in a line. The person standing at the i-th position from the front is the person Pi. For Q queries given as input, we have to figure out which of the two indexes is closer from the front.

Some Language difficulties

I must admit I got confused by the English translation, for example the person standing further to the front, does it mean closer from the front or further away to the back? Good thing that examples are provided! Likewise the Between the person Ai and Bi, I didn’t know if we have to compare all the indexes between Ai and Bi, or just Ai and Bi. I guess one small conclusion in terms of time management is to read the examples straight away, without beginning to think about a solution, as I can waste time going in the wrong direction.

The algorithm

It is pretty simple - I store people’s positions in a HashMap and for each query will compare the two entries. So the code becomes:

input! {
       n: usize,
       p: [usize; n],
       q: usize,
       queries: [(usize, usize); q],
    };

    let mut positions: HashMap<usize, usize> = HashMap::new();
    for i in 0..n {
        positions.insert(p[i], i);
    }

    for (a,b) in queries {
        let pa = positions.get(&a).unwrap();
        let pb = positions.get(&b).unwrap();
       if pa < pb {
        println!("{a}");
       } else {
        println!("{b}");
       }
    }

Problem C - Many Replacement

I spent slightly too much time on this problem, maybe trying to optimize too early. Here is the problem:

You are given a string S of length N consisting of lowercase English letters. You will perform an operation Q times on the string S. The i-th operation (1≤i≤Q) is represented by a pair of characters (c_i, d_i), which corresponds to the following operation:

Replace all occurrences of the character c_i in S with the character d_i. Print the string S after all operations are completed.

Solution

The string is long enough (2E5) that I don’t think we will have time to actually implement all the changes (i.e. for each query, look at all the characters and change those). I thought about simplifying the queries, by associativity: If all the a are changed into b, and then the b in c, then it is equivalent that all the a are changed into c. Or so that was the initial thought:

• Keep track of all the transformations for all characters, to get the “equivalent” query. (So O(Q) complexity).
• Then at the end we can just apply the “summary” query, we will have then a O(n) complexity.

The hole in this logic is that if all a are changed into b and then b into c, then we indeed have all a into c, but as well the b into c. (Because some b might be in the string from the beginning).

One way to think about it is the following:

• I will apply successive queries in the very special case that S0="abcdefghijklmnopqrstuvwxyz"
• Worst case scenario is that I look for every query 26 letters -> complexity of 26*Q = O(Q)
• Then I will apply this filter to each character of S, i.e. if the character is e then I will see which letter is contained in the 5th element of Q(S0)

The code is then:

fn main() {
    input! {
       n: usize,
       s: Chars,
       q: usize,
       queries: [(char, char);q],
    };
    // we store the final output of each character into a Map
    let mut outputs: HashMap<char, char> = HashMap::new();
    let mut dest = s.clone();

    for i in 'a'..='z' {
        outputs.insert(i, i);
    }

    for (a, b) in queries {
        // We apply each query to each letter to which this applies
        if outputs.get(&a).unwrap() == &a {
            outputs.insert(a, b);
        }
        for i in 'a'..='z' {
            if outputs.get(&i).unwrap() == &a {
                outputs.insert(i, b);
            }
        }
    }
    // We apply the summarized query to the initial string S
    for i in 0..n {
        dest[i] = *outputs.get(&dest[i]).unwrap();
    }
    println!("{}", dest.iter().collect::<String>());
}

Task D - Square Pair

Outcome

I understood the algorithm pretty fast, I made a mistake regarding how to apply this in the case of zero and spent too much time figuring this out. So close and yet so far!

Problem

You are given a sequence of non-negative integers $A=(A1,…,AN) of length N. Find the number of pairs of integers (i,j) that satisfy both of the following conditions:

• 1≤i<j≤N
• AiAj is a square number. Here, a non-negative integer a is called a square number when it can be expressed as a=d^2 using some non-negative integer d.

Solution

N is pretty big (2E5), as are each of the Ai, which means that the brute force algorithm will not work. In a first step, we will look at the decomposition of each Ai into prime factors. We can see that each Ai can be decomposed into squares of primes multiplied by some kind of residual. The idea is to notice that we can “simplify” each Ai by removing all the terms which are perfect squares, in other terms all the even powers of the prime factorization. And then to notice that AiAj is a perfect square if and only if Simplify(Ai) == Simplify(Aj). (And let’s suppose that both Ai and Aj are non zero now).

One first version of our algorithm would be:

• Simplify all Ai
• For each i, find the j > i such that Simplified(Ai) == Simplified(Aj)

The issue with the above algorithm is that it is of O(n^2) complexity, so we’ll get a bunch of TLEs. So we have to simplify the problem even further:

• Simplify all Ai
• Browse all Ai and “count” the occurences of the residuals.
• Each time there is a residual with count > 1, we know how many pairs can generate perfect squares: count*(count-1)/2. There we go, we have an algorithm in O(n) complexity. We just need to take care of the case then when Ai is equal to zero. Each Ai equal to zero, can make perfect square of any other number, so we have count_zero*(n - count_zero) to count for the pairs with one zero and one non-zero number, plus all the pairs with two zeros count_zero*(count_Zero-1)/2.

The algorithm is then:

fn get_prime_list(n: usize) -> Vec<usize> {
    let mut primes = vec![];
    let mut is_prime = vec![true; n + 1];
    is_prime[0] = false;
    is_prime[1] = false;
    for i in 2..n + 1 {
        if is_prime[i] {
            primes.push(i);
            let mut j = 2 * i;
            while j <= n {
                is_prime[j] = false;
                j += i;
            }
        }
    }
    primes
}

fn remove_squares(x: usize, prime_list: &Vec<usize>) -> usize {
    let mut x = x;
    for p in prime_list.iter() {
        if p * p > x {
            break;
        }
        while x % (p * p) == 0 {
            x /= p * p;
        }
    }
    x
}

fn main() {
    input! {
       n: usize,
       mut a: [usize; n],
    }
    let prime_list = get_prime_list(200_000);
    // We simplify each Ai by removing the squares
    for i in 0..n {
        if a[i] == 0 {
            continue;
        }
        a[i] = remove_squares(a[i], &prime_list);
    }
    let mut count_map = HashMap::new();
    let mut count_zero = 0;
    // We get count how many times each residual appear
    for i in 0..n {
        if a[i] == 0 {
            count_zero += 1;
            continue;
        }
        count_map.entry(a[i]).and_modify(|x|  *x = *x + 1).or_insert(1);
    }
    // We count the number of permutations for each residual
    let mut count = 0;
    for (_, v) in count_map.iter() {
        count += (v * (v - 1)) / 2;
    }
    // We look at the special case for Ai = 0
    if count_zero > 0 {
        count += count_zero* (n-count_zero) + (count_zero * (count_zero - 1)) / 2;
    }
    println!("{}", count ) ;
}