Equal sum grid partition ii

Array Hash Table Matrix Enumeration Prefix Sum

21.7% acceptance

Feb 25, 2026

Given an m x n matrix of positive integers, determine if a horizontal or vertical cut exists such that both sections are non-empty and have equal sum, OR can be made equal by discounting at most one cell (the rest of the section must remain connected).

Solution

Rust

Time O(n²)

Space O(n)

LeetCode

solution.rs

use std::collections::HashMap;

impl Solution {
  pub fn can_partition_grid(grid: Vec<Vec<i32>>) -> bool {
    let m = grid.len();
    let n = grid[0].len();
    let total: i64 = grid.iter().flat_map(|r| r.iter()).map(|&x| x as i64).sum();

    // value -> sorted unique rows that contain that value (used for horizontal cuts)
    let mut val_to_rows: HashMap<i32, Vec<usize>> = HashMap::new();
    // value -> sorted unique cols that contain that value (used for vertical cuts)
    let mut val_to_cols: HashMap<i32, Vec<usize>> = HashMap::new();
    for r in 0..m {
      for c in 0..n {
        val_to_rows.entry(grid[r][c]).or_default().push(r);
        val_to_cols.entry(grid[r][c]).or_default().push(c);
      }
    }
    for v in val_to_rows.values_mut() { v.sort_unstable(); v.dedup(); }
    for v in val_to_cols.values_mut() { v.sort_unstable(); v.dedup(); }

    // Does `target` appear in any row within [row_a, row_b]?
    let val_in_rows = |target: i32, row_a: usize, row_b: usize| -> bool {
      val_to_rows.get(&target).map_or(false, |rows| {
        let pos = rows.partition_point(|&r| r < row_a);
        pos < rows.len() && rows[pos] <= row_b
      })
    };

    // Does `target` appear in any col within [col_a, col_b]?
    let val_in_cols = |target: i32, col_a: usize, col_b: usize| -> bool {
      val_to_cols.get(&target).map_or(false, |cols| {
        let pos = cols.partition_point(|&c| c < col_a);
        pos < cols.len() && cols[pos] <= col_b
      })
    };

    // ── Horizontal cuts (cut after row r) ────────────────────────────────────
    // Top section : rows 0..=r  (h_top = r+1, width = n)
    // Bottom section: rows r+1..=m-1  (h_bot = m-r-1, width = n)
    let mut top_sum = 0i64;
    for r in 0..m - 1 {
      top_sum += grid[r].iter().map(|&x| x as i64).sum::<i64>();
      let bot_sum = total - top_sum;
      let diff = (top_sum - bot_sum).abs();
      if diff == 0 { return true; }
      if diff > i32::MAX as i64 { continue; }
      let target = diff as i32;

      let (row_a, row_b) = if top_sum > bot_sum { (0, r) } else { (r + 1, m - 1) };
      let h = row_b - row_a + 1;

      let found = if h >= 2 && n >= 2 {
        // Any cell is removable – just check existence in row range
        val_in_rows(target, row_a, row_b)
      } else if h == 1 && n >= 2 {
        // Single-row section: only the two end columns are removable
        grid[row_a][0] == target || grid[row_a][n - 1] == target
      } else if h >= 2 && n == 1 {
        // Single-column section: only the two end rows are removable
        grid[row_a][0] == target || grid[row_b][0] == target
      } else {
        false // 1×1 section
      };
      if found { return true; }
    }

    // ── Vertical cuts (cut after col c) ──────────────────────────────────────
    // Left section : cols 0..=c  (height = m, w_left = c+1)
    // Right section: cols c+1..=n-1  (height = m, w_right = n-c-1)
    let mut left_sum = 0i64;
    for c in 0..n - 1 {
      left_sum += grid.iter().map(|row| row[c] as i64).sum::<i64>();
      let right_sum = total - left_sum;
      let diff = (left_sum - right_sum).abs();
      if diff == 0 { return true; }
      if diff > i32::MAX as i64 { continue; }
      let target = diff as i32;

      let (col_a, col_b) = if left_sum > right_sum { (0, c) } else { (c + 1, n - 1) };
      let w = col_b - col_a + 1;

      let found = if m >= 2 && w >= 2 {
        // Any cell is removable – just check existence in col range
        val_in_cols(target, col_a, col_b)
      } else if m == 1 && w >= 2 {
        // Single-row grid: only the two end columns of the section are removable
        grid[0][col_a] == target || grid[0][col_b] == target
      } else if m >= 2 && w == 1 {
        // Single-column section: only the two end rows are removable
        grid[0][col_a] == target || grid[m - 1][col_a] == target
      } else {
        false // 1×1 section
      };
      if found { return true; }
    }

    false
  }
}