Problem Description

You need to design a class called NeighborSum that works with a square grid (n×n 2D array) containing distinct elements ranging from 0 to n²-1.

The class should support three operations:

1. Constructor NeighborSum(int[][] grid): Initializes the object with the given grid.

2. Adjacent Sum adjacentSum(int value): Returns the sum of elements that are directly adjacent to the given value in the grid. Adjacent means the cells that share an edge with the cell containing value - specifically the cells directly above, below, left, and right.

3. Diagonal Sum diagonalSum(int value): Returns the sum of elements that are diagonally adjacent to the given value in the grid. Diagonally adjacent means the cells at the four corners relative to the cell containing value - specifically the cells at top-left, top-right, bottom-left, and bottom-right positions.

The solution uses a hash table to store the position (row, column) of each element in the grid for quick lookup. When calculating sums, it uses direction vectors to iterate through the required neighboring cells. The dirs tuple contains two sets of direction offsets - one for adjacent cells (-1, 0, 1, 0, -1) and another for diagonal cells (-1, 1, 1, -1, -1). The pairwise function is used to extract consecutive pairs of values that represent the row and column offsets for each neighbor. The implementation checks boundaries to ensure only valid grid positions are included in the sum.

Intuition

The key insight is that we need to frequently find where a specific value is located in the grid and then calculate sums of its neighbors. Since values are distinct and we might query the same value multiple times, preprocessing the grid to create a value-to-position mapping makes sense.

Instead of searching through the entire grid every time we need to find a value (which would take O(n²) time), we can build a hash table during initialization that maps each value to its coordinates (i, j). This gives us O(1) lookup time for any value.

For calculating neighbor sums, we observe that:

• Adjacent neighbors are always at a fixed offset: up (-1, 0), right (0, 1), down (1, 0), left (0, -1)
• Diagonal neighbors are also at fixed offsets: top-left (-1, -1), top-right (-1, 1), bottom-right (1, 1), bottom-left (1, -1)

We can encode these directions cleverly. Notice that adjacent directions can be represented as consecutive pairs from the sequence (-1, 0, 1, 0, -1). Using pairwise, we get (-1, 0), (0, 1), (1, 0), (0, -1). Similarly, diagonal directions come from (-1, 1, 1, -1, -1) giving us (-1, 1), (1, 1), (1, -1), (-1, -1).

This unified approach allows us to write a single helper function cal that handles both types of sums by just switching the direction array. We simply iterate through the appropriate neighbors, check if they're within bounds, and accumulate their values.

Solution Approach

The solution uses a hash table to efficiently map values to their positions in the grid. Here's how the implementation works:

Initialization (__init__ method):

• Store the grid reference in self.grid
• Create an empty hash table self.d to store value-to-position mappings
• Define direction vectors in self.dirs tuple:
  • First element (-1, 0, 1, 0, -1) encodes adjacent directions
  • Second element (-1, 1, 1, -1, -1) encodes diagonal directions
• Iterate through the entire grid using enumerate to get both indices and values
• For each element x at position (i, j), store self.d[x] = (i, j)

Helper Function (cal method): This unified function handles both adjacent and diagonal sum calculations:

• Retrieve the position (i, j) of the given value from the hash table self.d[value]
• Initialize sum s = 0
• Select the appropriate direction array using index k (0 for adjacent, 1 for diagonal)
• Use pairwise(self.dirs[k]) to generate consecutive pairs representing direction offsets:
  • For adjacent (k=0): generates (-1, 0), (0, 1), (1, 0), (0, -1)
  • For diagonal (k=1): generates (-1, 1), (1, 1), (1, -1), (-1, -1)
• For each direction offset (a, b):
  • Calculate neighbor position: x, y = i + a, j + b
  • Check if the position is within grid boundaries: 0 <= x < len(self.grid) and 0 <= y < len(self.grid[0])
  • If valid, add the neighbor's value to the sum: s += self.grid[x][y]
• Return the calculated sum

Public Methods:

• adjacentSum(value): Simply calls self.cal(value, 0) to calculate adjacent sum
• diagonalSum(value): Simply calls self.cal(value, 1) to calculate diagonal sum

The time complexity for initialization is O(n²) where n is the grid size, while both sum operations run in O(1) time since we're checking at most 4 neighbors and using O(1) hash table lookup.

Example Walkthrough

Let's walk through a concrete example with a 3×3 grid:

grid = [[0, 1, 2],
        [3, 4, 5],
        [6, 7, 8]]

Step 1: Initialization When we create NeighborSum(grid), the constructor builds a hash table mapping each value to its position:

• d[0] = (0, 0), d[1] = (0, 1), d[2] = (0, 2)
• d[3] = (1, 0), d[4] = (1, 1), d[5] = (1, 2)
• d[6] = (2, 0), d[7] = (2, 1), d[8] = (2, 2)

The direction vectors are also stored:

• dirs[0] = (-1, 0, 1, 0, -1) for adjacent neighbors
• dirs[1] = (-1, 1, 1, -1, -1) for diagonal neighbors

Step 2: Calculate adjacentSum(4)

1. Look up position of value 4: d[4] = (1, 1) (center of grid)
2. Use pairwise(dirs[0]) to get adjacent offsets:
   • (-1, 0) → position (0, 1) → value = 1 ✓
   • (0, 1) → position (1, 2) → value = 5 ✓
   • (1, 0) → position (2, 1) → value = 7 ✓
   • (0, -1) → position (1, 0) → value = 3 ✓
3. Sum = 1 + 5 + 7 + 3 = 16

Step 3: Calculate diagonalSum(4)

1. Look up position of value 4: d[4] = (1, 1)
2. Use pairwise(dirs[1]) to get diagonal offsets:
   • (-1, 1) → position (0, 2) → value = 2 ✓
   • (1, 1) → position (2, 2) → value = 8 ✓
   • (1, -1) → position (2, 0) → value = 6 ✓
   • (-1, -1) → position (0, 0) → value = 0 ✓
3. Sum = 2 + 8 + 6 + 0 = 16

Step 4: Calculate adjacentSum(0) (corner element)

1. Look up position of value 0: d[0] = (0, 0) (top-left corner)
2. Check each adjacent position:
   • (-1, 0) → position (-1, 0) → out of bounds ✗
   • (0, 1) → position (0, 1) → value = 1 ✓
   • (1, 0) → position (1, 0) → value = 3 ✓
   • (0, -1) → position (0, -1) → out of bounds ✗
3. Sum = 1 + 3 = 4

The boundary checking ensures we only sum valid neighbors, making the solution robust for edge and corner elements.

Solution Implementation

Time and Space Complexity

Time Complexity:

• Initialization (__init__): The constructor iterates through every element in the grid to build a dictionary mapping each value to its position (i, j). With a grid of size m × n, this requires visiting all elements once, resulting in O(m × n) time complexity.

• adjacentSum(value) and diagonalSum(value): Both methods call the cal function, which performs a constant number of operations. The cal function uses pairwise to iterate through exactly 4 directions (either adjacent or diagonal), and for each direction, it performs constant-time operations (boundary checking and grid access). The position lookup self.d[value] is O(1) due to the hash table. Therefore, both methods have O(1) time complexity.

Space Complexity:

• The class stores the entire grid and creates a dictionary d that contains an entry for each element in the grid. The dictionary stores m × n key-value pairs where each value is a tuple of coordinates. Therefore, the space complexity is O(m × n).

• The dirs tuple contains a fixed number of elements (10 total across both tuples), which is O(1) additional space.

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Incorrect Direction Vector Encoding

One of the most common mistakes is misunderstanding how the direction vectors are encoded using the compact form (-1, 0, 1, 0, -1) and how pairwise extracts the direction offsets.

The Pitfall: Developers might incorrectly assume the direction array directly represents individual offsets or try to iterate through it without using pairwise, leading to incorrect neighbor calculations.

# WRONG: Trying to use directions directly
directions = (-1, 0, 1, 0, -1)
for offset in directions:  # This doesn't give you (dx, dy) pairs!
    neighbor_row = row + offset
    # Missing the column offset entirely

The Solution: Understanding that pairwise creates consecutive overlapping pairs from the sequence. For (-1, 0, 1, 0, -1), it generates: (-1, 0), (0, 1), (1, 0), (0, -1) which represent (row_offset, col_offset) for up, right, down, and left respectively.

# CORRECT: Using pairwise to extract direction pairs
for dx, dy in pairwise(directions[0]):
    neighbor_row = row + dx
    neighbor_col = col + dy

2. Boundary Checking Logic Errors

Another frequent mistake involves incorrect boundary validation, especially forgetting that the grid might not be perfectly square or using wrong comparison operators.

The Pitfall:

# WRONG: Assuming square grid and using wrong bounds
if 0 <= neighbor_row <= len(self.grid) and 0 <= neighbor_col <= len(self.grid):
    # This allows neighbor_row == len(self.grid), which is out of bounds!

The Solution: Always use strict less-than for upper bounds and check both dimensions separately:

# CORRECT: Proper boundary checking
if (0 <= neighbor_row < len(self.grid) and 
    0 <= neighbor_col < len(self.grid[0])):
    total_sum += self.grid[neighbor_row][neighbor_col]

3. Missing Import or Python Version Incompatibility

The pairwise function is only available in Python 3.10+. Using it in earlier versions will cause an ImportError.

The Pitfall:

from itertools import pairwise  # ImportError in Python < 3.10

The Solution: For Python versions before 3.10, implement your own pairwise function or use an alternative approach:

# Alternative 1: Manual implementation of pairwise
def pairwise(iterable):
    a = iter(iterable)
    b = iter(iterable)
    next(b, None)
    return zip(a, b)

# Alternative 2: Use explicit direction arrays
class NeighborSum:
    def __init__(self, grid):
        self.grid = grid
        self.value_to_position = {}
        # More explicit but clearer
        self.adjacent_dirs = [(-1, 0), (0, 1), (1, 0), (0, -1)]
        self.diagonal_dirs = [(-1, -1), (-1, 1), (1, 1), (1, -1)]
      
        for i, row in enumerate(grid):
            for j, value in enumerate(row):
                self.value_to_position[value] = (i, j)
  
    def _calculate_neighbor_sum(self, value, use_diagonal):
        row, col = self.value_to_position[value]
        directions = self.diagonal_dirs if use_diagonal else self.adjacent_dirs
        total = 0
      
        for dx, dy in directions:
            nr, nc = row + dx, col + dy
            if 0 <= nr < len(self.grid) and 0 <= nc < len(self.grid[0]):
                total += self.grid[nr][nc]
      
        return total

This alternative approach is more readable and works across all Python versions, making it a safer choice for production code.