Write given number as sums of five fifth powers

fivesums.rs

use rayon::prelude::*;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::thread;
use std::time::{Duration, Instant};
use std::io::{self, Write};

/// Finds sets of 3 numbers within `[range_start, range_end]`
/// such that the sum of their 5th powers equals `target_sum`.
/// Complexity: O(N^2 log N) time, O(N^2) space (compact)
pub fn find_three_sums(range_start: i64, range_end: i64, target_sum: i64) -> Vec<[i64; 3]> {
    let powers: Vec<i64> = (range_start..=range_end).map(|x| x.pow(5)).collect();
    let val_to_num: Vec<i64> = (range_start..=range_end).collect();
    let n = powers.len();

    if n > 65535 {
        panic!("Range too large for u16 index packing optimization (max 65535 items)");
    }

    // Optimization: Use packed indices (u32) instead of HashMap
    // This eliminates the massive memory spike caused by HashMap<i64, Vec<[usize; 2]>>
    // Memory usage: N^2/2 * 4 bytes. For N=10000, ~200MB.
    let mut pairs: Vec<u32> = Vec::with_capacity(n * n / 2);
    for b in 0..n {
        for c in b..n {
            // Pack b and c into u32
            let packed = ((b as u32) << 16) | (c as u32);
            pairs.push(packed);
        }
    }

    // Helper to compute sum
    let get_sum = |packed: u32| -> i64 {
        let b = (packed >> 16) as usize;
        let c = (packed & 0xFFFF) as usize;
        powers[b].saturating_add(powers[c])
    };

    // Sort by sum
    pairs.par_sort_unstable_by(|&p1, &p2| {
        let s1 = get_sum(p1);
        let s2 = get_sum(p2);
        s1.cmp(&s2)
    });

    // Search
    let results: Vec<[i64; 3]> = (0..n).into_par_iter().flat_map(|c| {
        let mut local_res = Vec::new();
        let needed = target_sum.saturating_sub(powers[c]);
        
        // Local helper for the closure
        let get_sum_local = |packed: u32| -> i64 {
            let b = (packed >> 16) as usize;
            let c = (packed & 0xFFFF) as usize;
            powers[b].saturating_add(powers[c])
        };

        // Binary search for the needed sum
        let search_res = pairs.binary_search_by(|&p| {
            get_sum_local(p).cmp(&needed)
        });

        if let Ok(idx) = search_res {
            // Found at least one match, expand to find all duplicates
            let mut left = idx;
            while left > 0 && get_sum_local(pairs[left - 1]) == needed {
                left -= 1;
            }
            let mut right = idx;
            while right < pairs.len() - 1 && get_sum_local(pairs[right + 1]) == needed {
                right += 1;
            }
            
            for i in left..=right {
                let packed = pairs[i];
                let a = (packed >> 16) as usize;
                let b = (packed & 0xFFFF) as usize;
                
                // Enforce ordering a <= b <= c
                if b <= c {
                    local_res.push([val_to_num[a], val_to_num[b], val_to_num[c]]);
                }
            }
        }
        local_res
    }).collect();

    results
}

/// Finds sets of 5 numbers within `[range_start, range_end]`
/// such that the sum of their 5th powers equals `target_sum`.
///
/// EXCLUDES any solution containing a cancelling pair (x, -x).
/// Optimized with sorted array + two-pointer search (O(N^3) time, O(N^2) space)
/// Memory Optimized: Uses packed indices (u32) instead of full structs to save 75% memory.
pub fn find_pure_five_sums(range_start: i64, range_end: i64, target_sum: i64) -> Vec<[i64; 5]> {
    let powers: Vec<i64> = (range_start..=range_end).map(|x| x.pow(5)).collect();
    let val_to_num: Vec<i64> = (range_start..=range_end).collect();
    let n = powers.len();

    if n > 65535 {
        panic!("Range too large for u16 index packing optimization (max 65535 items)");
    }

    // Optimization: Use packed indices (u32) to store (b, c).
    // High 16 bits = b, Low 16 bits = c.
    // We DO NOT store the sum. We recompute it on the fly.
    // This reduces memory usage from ~16 bytes/item to 4 bytes/item.
    let mut pairs: Vec<u32> = Vec::with_capacity(n * n / 2);
    for b in 0..n {
        for c in b..n {
            let vb = val_to_num[b];
            let vc = val_to_num[c];
            // Skip cancelling pairs like (-5, 5). Allow (0, 0).
            if vb != 0 && vb == -vc {
                continue;
            }
            // Pack b and c into u32
            let packed = ((b as u32) << 16) | (c as u32);
            pairs.push(packed);
        }
    }

    // Helper closure to compute sum from packed index
    let get_sum = |packed: u32| -> i64 {
        let b = (packed >> 16) as usize;
        let c = (packed & 0xFFFF) as usize;
        // Unsafe access could be faster but let's stick to safe Rust for now.
        // The compiler should optimize bounds checks since b, c < n is invariant.
        powers[b].saturating_add(powers[c])
    };

    // Sort by sum for two-pointer approach
    // We use par_sort_unstable_by to sort based on computed sum
    pairs.par_sort_unstable_by(|&p1, &p2| {
        let s1 = get_sum(p1);
        let s2 = get_sum(p2);
        s1.cmp(&s2)
    });

    let pairs = Arc::new(pairs);
    let val_to_num = Arc::new(val_to_num);
    let powers = Arc::new(powers);

    // Progress bar setup
    let counter = Arc::new(AtomicUsize::new(0));
    let total = n;
    let counter_clone = counter.clone();
    
    // Spawn progress thread
    let progress_handle = thread::spawn(move || {
        let start_time = Instant::now();
        let mut last_print = Instant::now();
        // Always show progress for debugging/visibility
        let show_threshold = Duration::from_millis(0);
        
        // Initial print
        print_progress(0, total);

        loop {
            thread::sleep(Duration::from_millis(50));
            let c = counter_clone.load(Ordering::Relaxed);
            
            if last_print.elapsed().as_millis() > 50 {
                print_progress(c, total);
                last_print = Instant::now();
            }

            if c >= total {
                break;
            }
        }
        
        // Final 100% print and newline
        print_progress(total, total);
        println!(); // Newline instead of clearing
    });

    let results: Vec<[i64; 5]> = (0..n).into_par_iter().flat_map(|a| {
        let mut local_results = Vec::new();
        let target_minus_a = target_sum.saturating_sub(powers[a]);
        
        let p_len = pairs.len();
        if p_len == 0 {
            counter.fetch_add(1, Ordering::Relaxed);
            return local_results;
        }

        let mut left = 0;
        let mut right = p_len - 1;

        // Local helper to avoid repeated closure creation
        let get_sum_local = |idx: usize| -> i64 {
            let packed = pairs[idx];
            let b = (packed >> 16) as usize;
            let c = (packed & 0xFFFF) as usize;
            powers[b].saturating_add(powers[c])
        };

        while left < p_len && right < p_len && left <= right {
             let sum_l = get_sum_local(left);
             let sum_r = get_sum_local(right);
             let sum = sum_l.saturating_add(sum_r);
             
             if sum < target_minus_a {
                 left += 1;
             } else if sum > target_minus_a {
                 if right == 0 { break; }
                 right -= 1;
             } else {
                 // Found a match!
                 // Handle duplicates.
                 let l_start = left;
                 let mut l_end = left;
                 // Look ahead for same sum
                 while l_end < p_len && get_sum_local(l_end) == sum_l {
                     l_end += 1;
                 }
                 
                 let r_end = right; 
                 let mut r_start = right;
                 // Look back for same sum
                 while r_start > 0 && get_sum_local(r_start - 1) == sum_r {
                     r_start -= 1;
                 }
                 
                 // Cartesian product of these two blocks
                 for i in l_start..l_end {
                     for j in r_start..=r_end {
                         // Check constraints: a <= b <= c <= d <= e
                         let packed_i = pairs[i];
                         let packed_j = pairs[j];
                         
                         let b_idx = (packed_i >> 16) as usize;
                         let c_idx = (packed_i & 0xFFFF) as usize;
                         let d_idx = (packed_j >> 16) as usize;
                         let e_idx = (packed_j & 0xFFFF) as usize;
                         
                         if a <= b_idx && c_idx <= d_idx {
                             let nums = [
                                 val_to_num[a],
                                 val_to_num[b_idx],
                                 val_to_num[c_idx],
                                 val_to_num[d_idx],
                                 val_to_num[e_idx]
                             ];
                             
                             if !has_cancelling_pair(&nums) {
                                 local_results.push(nums);
                             }
                         }
                     }
                 }
                 
                 // Advance pointers
                 left = l_end;
                 if r_start == 0 { break; }
                 right = r_start - 1;
             }
        }

        counter.fetch_add(1, Ordering::Relaxed);
        local_results
    }).collect();

    progress_handle.join().unwrap();
    results
}

fn print_progress(current: usize, total: usize) {
    let width = 40;
    let percentage = current as f64 / total as f64;
    let filled = (percentage * width as f64) as usize;
    let empty = width - filled;
    let bar: String = "=".repeat(filled) + &" ".repeat(empty);
    print!("\rProgress: [{}] {:.1}% ({}/{})", bar, percentage * 100.0, current, total);
    io::stdout().flush().unwrap();
}

fn has_cancelling_pair(nums: &[i64; 5]) -> bool {
    for i in 0..nums.len() {
        for j in i+1..nums.len() {
            if nums[i] != 0 && nums[i] == -nums[j] {
                return true;
            }
        }
    }
    false
}

fn main() {
    // Example usage:
    // Find 5 numbers in range [0, 100] whose 5th powers sum to a specific number.
    // Example: 1^5 + 2^5 + 3^5 + 4^5 + 5^5 = 1 + 32 + 243 + 1024 + 3125 = 4425

    let start = -5000;
    let end = 5000;
    for target in 1919810..=1919810 {
        fun_name(start, end, target);
    }
}

fn fun_name(start: i64, end: i64, target: i64) {
    println!(
        "range [{}, {}] summing to {}...",
        start, end, target
    );

    // Step 1: Find 3-sum solutions (which imply infinite 5-sum solutions with cancelling pairs)
    let three_sums = find_three_sums(start, end, target);
    if !three_sums.is_empty() {
        println!("    Found solutions with cancelling pairs (x, -x):");
        for s in &three_sums {
            println!("      Base {:?} + pairs (+/- x)", s);
        }
    }

    // Step 2: Find pure 5-sum solutions
    let solutions = find_pure_five_sums(start, end, target);

    if solutions.is_empty() && three_sums.is_empty() {
        println!("    No solutions found.");
    } else {
        for s in solutions {
            println!(
                "    Found pure 5-sum: {:?} -> {}^{} + {}^{} + {}^{} + {}^{} + {}^{} = {}",
                s, s[0], 5, s[1], 5, s[2], 5, s[3], 5, s[4], 5, target
            );
        }
    }
}