LeetCode 1808. Maximize Number of Nice Divisors

Problem

leetcode.com/problems/maximize-number-of-nice-divisors/

Maximize Number of Nice Divisors - LeetCode

Solving

The recurrence relation is:

f(k) = k when k <= 4
f(k) = 3 * f(k - 3) otherwise

Because the number of prime factors is at most 10^9, the equation can not simply be applied. (the time complexity is O(n) when n is the number of prime factors)

For a given k , we can rewrite k as:

k = 6 * x + r (0 <= r < 6)

Using the equation above, we can redefine the recurrence relation as:

f(1) = 1
f(2) = 2
f(3) = 3
f(4) = 4
f(5) = 3 * f(5-3) = 6
f(6) = 3 * f(6-3) = 9
f(7) = 3 * f(7-3) = 12
f(8) = 3 * f(8-3) = 18
f(k) = f(6x) * f(r)
     = f(3x) * f(3x) * f(r) 
     = f(3x)^2 * f(r) 

where k = 6 * x + r and 2 <= r < 6

Note that r should be at least 2, otherwise the equation would not reduce the optimal solution.

Because r < 6, we can calculate f(r) in O(1), and x >= 1.
So the entire time complexity turns into O(log n).

Solution

const M: i64 = 1000000007;

fn f(n: i64) -> i64 {
    match n {
        0 => 1,
        1 | 2 | 3 | 4 => n,
        5 => 6,
        6 => 9,
        7 => 12,
        8 => 18,
        _ => {
            let d = (n - 2) / 6 * 3;
            let x = f(d);
            let x = x * x % M;
            x * f(n - d * 2) % M    
        }
    }
}

impl Solution {
    pub fn max_nice_divisors(prime_factors: i32) -> i32 {
        f(prime_factors as i64) as i32
    }
}