Solution to Facebook's Hacker Cup 2014 first problem: Square Detector

01-square-detector.py

#!/usr/bin/env python
# encoding: utf-8
"""Square Detector

https://www.facebook.com/hackercup/problems.php?pid=318555664954399&round=598486203541358

You want to write an image detection system that is able to recognize different geometric shapes. In the first version of the system you settled with just being able to detect filled squares on a grid.

You are given a grid of NxN square cells. Each cell is either white or black. Your task is to detect whether all the black cells form a square shape.

Input
The first line of the input consists of a single number T, the number of test cases.

Each test case starts with a line containing a single integer N. Each of the subsequent N lines contain N characters. Each character is either "." symbolizing a white cell, or "#" symbolizing a black cell. Every test case contains at least one black cell.

Output
For each test case i numbered from 1 to T, output "Case #i: ", followed by YES or NO depending on whether or not all the black cells form a completely filled square with edges parallel to the grid of cells.
"""
from __future__ import division
import itertools
import numpy as np
import sys


def read_line():
    return sys.stdin.readline().strip()


def read_ints():
    return map(int, read_line().split())


def read_lines(n):
    while n > 0:
        yield read_line()
        n -= 1


def char_to_num(c):
    if c == '#':
        return 1
    return 0


def solve(lines):
    M = np.mat([map(char_to_num, l) for l in lines])
    total = np.sum(M)
    # quick check: number of black cells must be exact square
    sq = np.sqrt(total)
    if sq != int(sq):
        return False
    # find first black cell
    n, m = M.shape
    for i, j in itertools.product(xrange(n), xrange(m)):
        if M[i, j] == 1:
            break
    # get sub matrix
    S = M[i:i+sq, j:j+sq]
    # are all ones?
    return (S == 1).all()


def main():
    cases = read_ints()[0]
    for i in xrange(cases):
        n = read_ints()[0]
        result = solve(read_lines(n))
        if result:
            output = 'YES'
        else:
            output = 'NO'
        print 'Case #{}: {}'.format(i + 1, output)
        cases -= 1


if __name__ == '__main__':
    main()

sample.in

5
4
..##
..##
....
....
4
..##
..##
#...
....
4
####
#..#
#..#
####
5
#####
#####
#####
#####
.....
5
#####
#####
#####
#####
#####

sample.out

Case #1: YES
Case #2: NO
Case #3: NO
Case #4: NO
Case #5: YES