Skip to content

3548. Equal Sum Grid Partition II

  • Time: $O(mn)$
  • Space: $O(mn)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
class Solution {
 public:
  bool canPartitionGrid(vector<vector<int>>& grid) {
    const long sum = accumulate(grid.begin(), grid.end(), 0L,
                                [](long acc, const vector<int>& row) {
      return acc + accumulate(row.begin(), row.end(), 0L);
    });
    return canPartition(grid, sum) || canPartition(reversed(grid), sum) ||
           canPartition(reversed(transposed(grid)), sum) ||
           canPartition(transposed(grid), sum);
  }

 private:
  bool canPartition(const vector<vector<int>>& grid, long sum) {
    long topSum = 0;
    unordered_set<int> seen;
    for (int i = 0; i < grid.size(); ++i) {
      topSum += accumulate(grid[i].begin(), grid[i].end(), 0L);
      const long botSum = sum - topSum;
      seen.insert(grid[i].begin(), grid[i].end());
      if (topSum - botSum == 0 || topSum - botSum == grid[0][0] ||
          topSum - botSum == grid[0].back() || topSum - botSum == grid[i][0])
        return true;
      if (grid[0].size() > 1 && i > 0 && seen.count(topSum - botSum))
        return true;
    }
    return false;
  }

  vector<vector<int>> transposed(const vector<vector<int>>& grid) {
    vector<vector<int>> res(grid[0].size(), vector<int>(grid.size()));
    for (int i = 0; i < grid.size(); ++i)
      for (int j = 0; j < grid[0].size(); ++j)
        res[j][i] = grid[i][j];
    return res;
  }

  vector<vector<int>> reversed(const vector<vector<int>>& grid) {
    return {grid.rbegin(), grid.rend()};
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
class Solution {
  public boolean canPartitionGrid(int[][] grid) {
    final long sum = Arrays.stream(grid).flatMapToInt(Arrays::stream).asLongStream().sum();
    return canPartition(grid, sum) || canPartition(reversed(grid), sum) ||
        canPartition(reversed(transposed(grid)), sum) || canPartition(transposed(grid), sum);
  }

  private boolean canPartition(int[][] grid, long sum) {
    long topSum = 0;
    Set<Integer> seen = new HashSet<>();
    for (int i = 0; i < grid.length; ++i) {
      topSum += Arrays.stream(grid[i]).asLongStream().sum();
      final long botSum = sum - topSum;
      Arrays.stream(grid[i]).forEach(seen::add);
      if (topSum - botSum == 0 || topSum - botSum == grid[0][0] ||
          topSum - botSum == grid[0][grid[0].length - 1] || topSum - botSum == grid[i][0])
        return true;
      if (grid[0].length > 1 && i > 0 && seen.contains((int) (topSum - botSum)))
        return true;
    }
    return false;
  }

  private int[][] transposed(int[][] grid) {
    int[][] res = new int[grid[0].length][grid.length];
    for (int i = 0; i < grid.length; ++i)
      for (int j = 0; j < grid[0].length; ++j)
        res[j][i] = grid[i][j];
    return res;
  }

  private int[][] reversed(int[][] grid) {
    return Arrays.stream(grid).collect(Collectors.collectingAndThen(Collectors.toList(), list -> {
      Collections.reverse(list);
      return list.toArray(new int[0][]);
    }));
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
class Solution:
  def canPartitionGrid(self, grid: list[list[int]]) -> bool:
    summ = sum(map(sum, grid))

    def canPartition(grid: list[list[int]]) -> bool:
      topSum = 0
      seen = set()
      for i, row in enumerate(grid):
        topSum += sum(row)
        botSum = summ - topSum
        seen |= set(row)
        if topSum - botSum in (0, grid[0][0],  grid[0][-1], row[0]):
          return True
        if len(grid[0]) > 1 and i > 0 and topSum - botSum in seen:
          return True
      return False

    return (canPartition(grid) or
            canPartition(grid[::-1]) or
            canPartition(list(zip(*grid))[::-1]) or
            canPartition(list(zip(*grid))))