Maximum sized array

Binary Search Bit Manipulation

51.2% acceptance

Mar 31, 2026

Given a positive integer s, let A be a 3D array of dimensions n × n × n, where each element A[i][j][k] is defined as: A[i][j][k] = i * (j OR k), where 0 <= i, j, k < n. Return the maximum possible value of n such that the sum of all elements in array A does not exceed s.

Solution

Rust

Time O(n log n)

Space O(1)

LeetCode

solution.rs

impl Solution {
  pub fn max_sized_array(s: i64) -> i32 {
    // sum = sum_i(i) * sum_{j,k}(j|k) for j,k in [0,n)
    // sum_i(i) = n*(n-1)/2
    // We need to compute sum_{j,k in [0,n)} (j|k)
    // For each bit b, count pairs where bit b is set in j|k:
    // = n^2 - (count where bit b is 0 in both j and k)^2
    // For bit b with period 2^(b+1), count of values in [0,n) with bit b unset:
    // full = n / (2^(b+1)) * 2^b, partial = min(n % 2^(b+1), 2^b)
    // unset(b) = full + partial
    // jk_sum = sum over bits of (n^2 - unset(b)^2) * 2^b
    
    let compute_sum = |n: i64| -> i64 {
      let sum_i = n * (n - 1) / 2;
      let mut jk_sum: i64 = 0;
      let nn = n * n;
      for b in 0..20 {
        let period = 1i64 << (b + 1);
        let half = 1i64 << b;
        let full = (n / period) * half;
        let partial = std::cmp::min(n % period, half);
        let unset = full + partial;
        jk_sum += (nn - unset * unset) * half;
      }
      sum_i * jk_sum
    };
    
    let mut lo = 1i64;
    let mut hi = 1200i64; // upper bound estimate
    while lo < hi {
      let mid = (lo + hi + 1) / 2;
      if compute_sum(mid) <= s {
        lo = mid;
      } else {
        hi = mid - 1;
      }
    }
    lo as i32
  }
}