Biggest forest (Amazon interview question)

Viewed 5922

Each cell is either water 'W' or a tree 'T'. Given the information about the field, print the size of the largest forest. The size of a forest is the number of trees in it. See the sample case for clarity

INPUT:

First-line contains the size of the matrix N. The next N lines contain N characters each, either 'W' or 'T'.

OUTPUT:

Print the size of the biggest forest.

Sample input:

5
TTTWW
TWWTT
TWWTT
TWTTT
WWTTT

Expected Output: 10

My code:

t_cases = int(input())
k1 = 0
k2 = 0
for _ in range(t_cases):
    list1 = (input())
    z = 0
    list2 = []
    for i in range(len(list1)):
        z = list1.count('T')
        if list1[i] == "W":
            break
        elif list1[i] == "T":
            list2.append(list1[i])
            
    k1 = k1 + list2.count('T')
    if z > list2.count('T'):
        k2 = k2 + (z - list2.count('T'))
    else: 
        k2 = k2 + (list2.count('T')- z)
if k1 > k2:
    print(k1)
else: 
    print(k2)

My code satisfies the sample input but fails each and every test case. This code calculates the sum of tress before 'W' in all the cases and added them to k1. Similarly, k2 is the sum of all the trees after 'W'.

Note: Recursion could be also used!

5 Answers

This is essentially a classic flood-fill algorithm in disguise. For each tree that you see, you can run a flood-fill to find all the trees in the same forest, and from there you just need to return the maximum number of trees you find.

One way to do a flood-fill is with a breadth-first search. Here's some simple pseudocode for this; I'll leave the translation as an exercise since this is an interview practice problem.

max_forest = 0
for each location:
    if it's a tree, and you haven't visited it yet:
        max_forest = max(max_forest, size_of_forest(location))
return max_forest

size_of_forest(location):
    if this location has been visited already, return 0
    make a worklist of locations, initially just the start.
    
    size = 1
    while the worklist isn't empty:
        remove one element from the worklist.
        increment size.

        for each neighboring tree:
            if that location isn't yet visited:
                mark that location visited.
                increment size.
                add the location to the worklist.

        return size

Another approach would be to use a depth-first search. Here's some pseudocode:

size_of_forest(location):
    if this location is visited, return 0
    mark this location as visited

    result = 0
    for each neighboring tree:
        result += size_of_forest(that tree)

    return result

There are a bunch of questions you'll need to work through to turn this into working code. How will you track which locations have been visited? How will you iterate over the neighboring trees?

More abstractly, this problem is equivalent to finding the size of the largest connected component of the graph formed by having one node per tree, with edges between trees when they are adjacent to one another. The BFS and DFS pseudocode I've given here are the general BFS and DFS algorithms, just specialized for this particular case.

These two algorithms - BFS and DFS - are really good ones to know if you're out doing job interviews. They come up all the time in practice and are real workhorses once you know how to use them. (I've lost count of how many times I've needed to code these up!)

This is very easy to implement in e.g. python using a simple breath-first flood fill algorithm. Point being, you need to backtrack to make sure you properly search the forest. Try something like this

ex0="""
5
TTTWW
TWWTT
TWWTT
TWTTT
WWTTT
"""


def parse(data):
    """Returns the fields as a set of coordniates"""
    lines = iter(data.splitlines())
    next(lines)  # skip the size
    field = set()
    for y, line in enumerate(lines):
        for x, cell in enumerate(line):
            if cell == 'T':
                field.add((x, y))
    return field


DIRECTIONS = [(-1, 0), (1, 0), (0, -1), (0, 1)]


def neighbors(p):
    x, y = p
    for dx, dy in DIRECTIONS:
        yield x + dx, y + dy


def find_forests(field):
    unvisited = set(field)  # copy the field
    while unvisited:
        first = unvisited.pop()  # take any tree
        queue = [first]
        forest = set()
        while queue:
            p = queue.pop(0)
            if p in unvisited:
                unvisited.remove(p)
            forest.add(p)
            for n in neighbors(p):
                if n in unvisited:
                    queue.append(n)
        yield forest

forests = find_forests(parse(ex0))
print(max(map(len, forests)))  # find largest forest

Solution:

from collections import deque

n = int(input())
a = []
for i in range(n):
    a.append(input())

used = [[False for i in range(n)] for j in range(n)]
ans = 0
for i in range(n):
    for j in range(n):
        if used[i][j] or a[i][j] == 'W':
            continue
        q = deque()
        q.append((i, j))
        used[i][j] = True
        cnt = 0
        while q:
            (x, y) = q.pop()
            cnt += 1
            for dx, dy in [(1, 0), (-1, 0), (0, 1), (0, -1)]:
                x_ = x + dx
                y_ = y + dy
                if x_ >= 0 and x_ < n and y_ >= 0 and y_ < n:
                    if not used[x_][y_] and a[x_][y_] == 'T':
                        q.append((x_, y_))
                        used[x_][y_] = True
        ans = max(ans, cnt)
print(ans)

Explanation:

used - NxN array, where True means that cell is already visited. We go through all cells. If we find T and it is not visited before we start counting forest (cnt). For that we check right, left, top and botton cells. If they contains unvisited T then we add it to our forest. It is BFS algoritm.

Think about iterative or recursive methods to "explore" a forest - starting at a single tree and finding all adjacent tree, then finding adjacent trees to those. Then find a method of marking trees that have already been visited and skipping them as you iterate over the matrix.

Since no one mentioned it, here's an attempt at a labeling scheme that iterates sequentially on each row of the matrix and uses for additional space two rows and a dictionary of labels:

def f(m):
  n = len(m)
  
  # Stores final label and size
  # for each component
  labels = {}
  row1 = [0] * n
  row2 = [0] * n
  
  label = 0
        
  for i in range(n):
    for j in range(n):
      if m[i][j] == "T":
        if j == 0 or not row2[j-1]:
          # No label above or to the left
          if i == 0 or not row1[j]:
            label += 1
            labels[label] = {"size": 1, "label": label}
            row2[j] = label
            print("i, j: %s, %s; new label %s" % (i, j, label))
          # Label only above
          else:
            row2[j] = row1[j]
            labels[row1[j]]["size"] += 1
            print("i, j: %s, %s; continuing above label %s" % (i, j, labels[row1[j]]["label"]))
        # Label only to the left
        elif not row1[j]:
          row2[j] = row2[j-1]
          labels[row2[j-1]]["size"] += 1
          print("i, j: %s, %s; continuing left label %s" % (i, j, labels[row2[j-1]]["label"]))
        # Labels above and to the left
        else:
          row2[j] = row1[j]
          labels[row1[j]]["size"] += 1
          print("i, j: %s, %s; continuing above label %s" % (i, j, labels[row1[j]]["label"]))
          # Unequal labels above and to the left,
          # relabel
          if labels[row1[j]]["label"] != labels[row2[j-1]]["label"]:
            print("i, j: %s, %s; relabeling %s to %s" % (i, j, labels[row2[j-1]]["label"], labels[row1[j]]["label"]))
            labels[row2[j-1]]["label"] = labels[row1[j]]["label"]
    row1 = row2
    row2 = [0] * n
    to_del = set()
    for idx, lbl in enumerate(row1):
      if lbl and labels[lbl]["label"] != lbl:
        to_del.add(lbl)
        row1[idx] = labels[lbl]["label"]
    for lbl in to_del:
      target = labels[lbl]["label"]
      print("i: %s; combining label %s with label %s" % (i, lbl, target))
      labels[target]["size"] += labels[lbl]["size"]
      del labels[lbl]
        

  return labels
          
          
import sys

data = sys.stdin.readlines()
print(f(data))

stdin:

TTTWW
TWWTT
TWWTT
TWTTT
WWTTT

stdout:

i, j: 0, 0; new label 1
i, j: 0, 1; continuing left label 1
i, j: 0, 2; continuing left label 1
i, j: 1, 0; continuing above label 1
i, j: 1, 3; new label 2
i, j: 1, 4; continuing left label 2
i, j: 2, 0; continuing above label 1
i, j: 2, 3; continuing above label 2
i, j: 2, 4; continuing above label 2
i, j: 3, 0; continuing above label 1
i, j: 3, 2; new label 3
i, j: 3, 3; continuing above label 2
i, j: 3, 3; relabeling 3 to 2
i, j: 3, 4; continuing above label 2
i: 3; combining label 3 with label 2
i, j: 4, 2; continuing above label 2
i, j: 4, 3; continuing above label 2
i, j: 4, 4; continuing above label 2
{1: {'size': 6, 'label': 1}, 2: {'size': 10, 'label': 2}}

stdin:

WTTTT
WWWWT
WTTTT
WWTWT
WTTWW

stdout:

i, j: 0, 1; new label 1
i, j: 0, 2; continuing left label 1
i, j: 0, 3; continuing left label 1
i, j: 0, 4; continuing left label 1
i, j: 1, 4; continuing above label 1
i, j: 2, 1; new label 2
i, j: 2, 2; continuing left label 2
i, j: 2, 3; continuing left label 2
i, j: 2, 4; continuing above label 1
i, j: 2, 4; relabeling 2 to 1
i: 2; combining label 2 with label 1
i, j: 3, 2; continuing above label 1
i, j: 3, 4; continuing above label 1
i, j: 4, 1; new label 3
i, j: 4, 2; continuing above label 1
i, j: 4, 2; relabeling 3 to 1
i: 4; combining label 3 with label 1
{1: {'size': 13, 'label': 1}}
Related