meta hacker cup 2024, Problem E: Wildcard Submissions

meta_hacker_cup_2024_problem_e_wildcard_submissions.rs

/*
Problem E: Wildcard Submissions
https://www.facebook.com/codingcompetitions/hacker-cup/2024/round-1/problems/E
I didn't participate but I coded up the solution to E after reading over the contest analysis
*/

use rayon::prelude::*;
use std::io::{stdin, BufRead};

const MODULO: u64 = 998_244_353;

pub fn solve_e() {
    let mut it = stdin().lock().lines().map(|line| line.unwrap());
    let t: usize = it.next().unwrap().parse().unwrap();
    let mut test_cases = vec![];
    for case in 1..=t {
        let n: usize = it.next().unwrap().parse().unwrap();
        let mut words = vec![];
        for _ in 0..n {
            words.push(it.next().unwrap());
        }
        test_cases.push((case, words));
    }
    let lines: Vec<String> = test_cases
        .par_iter()
        .map(|(case, words)| {
            let ans = solve(words);
            eprintln!("Case #{case}: {ans}");
            format!("Case #{case}: {ans}")
        })
        .collect();
    println!("{}", lines.join("\n"));
}

fn solve(s: &Vec<String>) -> u64 {
    let lcp = get_lcp(s);
    let v = get_v(s);
    let n = s.len();
    let mut odd = 0;
    let mut even = 0;
    for s in 1..1 << n {
        let mut count = 0;
        for j in (0..lcp[s]).rev() {
            count += 1;
            if (v[s] >> j) & 1 == 1 {
                count *= 26;
            }
            count %= MODULO;
        }
        if s.count_ones() % 2 == 0 {
            even += count;
            even %= MODULO;
        } else {
            odd += count;
            odd %= MODULO;
        }
    }
    (1 + odd + MODULO - even) % MODULO
}

fn get_v(s: &Vec<String>) -> Vec<u128> {
    let n = s.len();
    let mut v = vec![0; 1 << n];
    for i in 0..n {
        for (j, c) in s[i].chars().enumerate() {
            if c == '?' {
                v[1 << i] |= 1 << j;
            }
        }
    }
    for s in 1usize..1 << n {
        if s.count_ones() > 1 {
            for i in 0..n {
                if (s >> i) & 1 == 1 {
                    v[s] = v[s - (1 << i)] & v[1 << i];
                    break;
                }
            }
        }
    }
    v
}

fn get_lcp(s: &Vec<String>) -> Vec<usize> {
    let n = s.len();
    let mut lcp = vec![0; 1 << n];
    lcp[0] = 1000;
    for i in 0..n {
        lcp[1 << i] = s[i].len();
        for j in 0..i {
            let common_count = s[i]
                .chars()
                .zip(s[j].chars())
                .take_while(|&(x, y)| x == '?' || y == '?' || x == y)
                .count();
            lcp[(1 << i) | (1 << j)] = common_count;
        }
    }
    for s in 1..1 << n {
        for i in 0..n {
            if (s >> i) & 1 == 1 {
                for j in i + 1..n {
                    if (s >> j) & 1 == 1 {
                        // todo: surely we don't need both i & j here. theres a simpler way
                        lcp[s] = lcp[s - (1 << i)]
                            .min(lcp[s - (1 << j)])
                            .min(lcp[(1 << i) | (1 << j)]);
                    }
                }
            }
        }
    }
    lcp
}