Number of ways to paint sheets

66.7% acceptance

Apr 3, 2026

You are given an integer n representing the number of sheets. You are also given an integer array limit of size m, where limit[i] is the maximum number of sheets that can be painted using color i. You must paint all n sheets under the following conditions: Exactly two distinct colors are used. Each color must cover a single contiguous segment of sheets. The number of sheets painted with color i cannot exceed limit[i]. Return an integer denoting the number of distinct ways to paint all sheets. Since the answer may be large, return it modulo 109 + 7. Note: Two ways differ if at least one sheet is painted with a different color.

Solution

Rust

Time O(n log n)

Space O(n)

LeetCode

solution.rs

impl Solution {
  pub fn number_of_ways(n: i32, limit: Vec<i32>) -> i32 {
    const MOD: i128 = 1_000_000_007;

    let max_segment = i64::from(n - 1);
    let mut caps: Vec<i64> = limit
      .into_iter()
      .map(|value| i64::from(value.min(n - 1)))
      .collect();
    caps.sort_unstable();

    let mut prefix = vec![0i128; caps.len() + 1];
    for index in 0..caps.len() {
      prefix[index + 1] = prefix[index] + i128::from(caps[index]);
    }

    let mut total = 0i128;
    for &cap in &caps {
      let need = max_segment - cap + 1;
      let start = caps.partition_point(|&other| other < need);
      let count = (caps.len() - start) as i128;
      let sum = prefix[caps.len()] - prefix[start];
      total += count * i128::from(cap) + sum - count * i128::from(max_segment);
    }

    for &cap in &caps {
      let diagonal = 2 * cap - max_segment;
      if diagonal > 0 {
        total -= i128::from(diagonal);
      }
    }

    total.rem_euclid(MOD) as i32
  }
}