Max connected area in 0/1 array -- Deep First Search

# 2017-11-17
# [email protected]

import random
import sys  # To print without newline.

def print_array(a, rows, cols, pos_cell='# ', zero_cell='. '):
  for r in xrange(rows):
    for c in xrange(cols):
      sys.stdout.write(pos_cell if a[r][c] else zero_cell)
    sys.stdout.write('\n')


def visit_and_count(a, rows, cols, visited, r, c, debug=False):
  """Visit the input array a started from (r, c) with visiting states in visited.

  visited array will be updated with the visited elements.

  Return the area of the connected area started at (r, c).
  """
  visited[r][c] = True
  if debug:
    sys.stdout.write('Visiting (%d, %d)\n' % (r, c))

  area = 1  # Current cell.
  dd = [(-1, 0), (0, 1), (1, 0), (0, -1)]
  for dr, dc in dd:
    next_row = r + dr
    next_col = c + dc
    # Only proceed with valid cell.
    if next_row < 0 or next_row >= rows or next_col < 0 or next_col >= cols:
      continue

    if a[next_row][next_col] == 1 and not visited[next_row][next_col]:
      area += visit_and_count(a, rows, cols, visited, next_row, next_col, debug)

  return area


def max_area_dfs(a, debug=False):
  """Find the max area of connected 1s in a.

  Returns: max_area, start_col, start_row.
  """

  rows = len(a)
  if not rows:
    return 0

  cols = len(a[0])

  visited = [[False] * cols for _ in xrange(rows)]
  max_area = 0
  # The coordiate of the start point of the max area.
  start_row, start_col = -1, -1

  for r in xrange(rows):
    for c in xrange(cols):
      if a[r][c] == 1 and not visited[r][c]:
        if debug:
          sys.stdout.write('Starting from (%d, %d)\n' % (r, c))
        area = visit_and_count(a, rows, cols, visited, r, c, debug)
        if area > max_area:
          max_area = area
          start_row = r
          start_col = c

  if debug:
    sys.stdout.write('Visited array.\n')
    print_array(visited, rows, cols)

  return max_area, start_row, start_col


def main():
  debug = False
  rows = 7
  cols = 7
  a = [[0] * cols for _ in xrange(rows)]
  # Randomly assign 1 to a.
  for r in xrange(rows):
    for c in xrange(cols):
      a[r][c] = 1 if random.random() < 0.4 else 0

  sys.stdout.write('Input array:\n')
  print_array(a, rows, cols)
  max_area, start_row, start_col = max_area_dfs(a, debug=debug)

  sys.stdout.write('Max area: %d\n' % max_area)
  sys.stdout.write('start_row: %d, start_col: %d\n' % (start_row, start_col))

  # Visualize the max area.
  visited = [[False] * cols for _ in xrange(rows)]
  visit_and_count(a, rows, cols, visited, start_row, start_col)
  print_array(visited, rows, cols, pos_cell='O ')


if __name__ == '__main__':
  main()