Code Jam 17/2/3

beaming_with_joy.py

from collections import defaultdict as dd
from itertools import product


# 2SAT linear solver START
def strongly_connected(graph):
    """ From Cormen: Strongly connected components chapter. """
    l = len(graph)
    normal = [list(graph[i]) for i in xrange(l)]
    reverse = [[] for _ in xrange(l)]
    for fr, tos in enumerate(normal):
        for to in tos:
            reverse[to].append(fr)

    finishing = [0 for _ in xrange(l)]
    was = [False] * l
    cf = 0
    for i in xrange(l):
        if was[i]:
            continue
        was[i] = True
        stack = [(i, 0)]
        while stack:
            curr, index = stack.pop()
            if index == len(normal[curr]):
                finishing[curr] = cf
                cf += 1
                continue
            stack.append((curr, index + 1))
            nxt = normal[curr][index]
            if not was[nxt]:
                was[nxt] = True
                stack.append((nxt, 0))

    order = sorted(range(l), key=lambda i: finishing[i], reverse=True)
    cc = 0
    components = [-1] * l
    was = [False] * l
    for el in order:
        if was[el]:
            continue
        was[el] = True
        stack = [el]
        components[el] = cc
        while stack:
            fr = stack.pop()
            for to in reverse[fr]:
                if was[to]:
                    continue
                components[to] = cc
                was[to] = True
                stack.append(to)

        cc += 1

    return components


def calc_strong_graph(graph, components):
    n = max(components) + 1
    ret = [[] for _ in xrange(n)]

    for fr, tos in enumerate(graph):
        for to in tos:
            cfr = components[fr]
            cto = components[to]
            if cfr != cto:
                ret[cfr].append(cto)

    return ret


def topological_order(graph):
    n = len(graph)
    need = [0] * n

    for tos in graph:
        for to in tos:
            need[to] += 1

    stack = [i for i in xrange(n) if need[i] == 0]
    order = 0
    ret = [None] * n
    while stack:
        fr = stack.pop()
        ret[fr] = order
        order += 1
        for to in graph[fr]:
            need[to] -= 1
            if need[to] == 0:
                stack.append(to)

    return ret


def two_sat_solve(conj_normal_form):
    """ Expects the input as a list of tuples
    each having 2 elements. If an element is
    i then we assume x_i and if it's -i
    we assume ~x_i. For example if we have
    (i, -j) it would mean x_i V ~x_j.

    The solution is a dict, mapping the variable
    index to True and False."""

    # Transform the variables
    fmap, rmap = {}, {}
    c = 0
    for els in conj_normal_form:
        for el in els:
            if abs(el) not in fmap:
                fmap[abs(el)] = c
                rmap[c] = abs(el)
                c += 1

    n = len(fmap)
    for key in rmap.keys():
        val = rmap[key]
        rmap[key + n] = -val
        fmap[-val] = key + n

    # Create the associated graph
    graph = [[] for _ in xrange(2 * n)]
    for i, j in conj_normal_form:
        for fr, to in ((-i, j), (-j, i)):
            graph[fmap[fr]].append(fmap[to])


    # Compute the connected components
    components = strongly_connected(graph)

    # Compute the reverse map from component to node
    rev_map = dd(set)
    for i, comp in enumerate(components):
        rev_map[comp].add(i)

    # Check if any variable is in the same component as its negative
    for values in rev_map.itervalues():
        for value in values:
            if ((value + n) % (2 * n)) in values:
                return False, {}

    # Calculate the gaph consisting of the strong components
    strong_graph = calc_strong_graph(graph, components)
    ns = len(strong_graph)

    # Get the topological order of the component graph
    top_order = topological_order(strong_graph)

    rev_order = sorted(range(ns), key=lambda i: top_order[i], reverse=True)
    values = [None] * 2 * n

    # Traverse the components in reverse topological order and
    # fill out the variables accordingly.
    for comp in rev_order:
        first = next(iter(rev_map[comp]))
        if values[first] is not None:
            continue
        other_comp = components[(first + n) % (2 * n)]
        for el in rev_map[comp]:
            values[el] = True
        for el in rev_map[other_comp]:
            values[el] = False

    # Calculate the return value
    ret = {}
    for i in xrange(n):
        ret[rmap[i]] = values[i]

    return True, ret
# 2SAT linear solver END


H, V = range(2)
L, R, U, D = range(4)


def show(case, possible, sol):
    if possible:
        print "Case #%s: POSSIBLE" % case
        for row in sol:
            print "".join(row)
    else:
        print "Case #%s: IMPOSSIBLE" % case


def is_in(x, y, r, c):
    return 0 <= x  and x < r and 0 <= y and y < c


def get_next(x, y, di):
    dxd = {
        L: (0, -1),
        R: (0, 1),
        U: (-1, 0),
        D: (1, 0)
    }
    dx, dy = dxd[di]
    return x + dx, y + dy


def turn(di, val):
    did = {
        L: {
            "/": D,
            "\\": U,
            ".": L,
        },
        R: {
            "/": U,
            "\\": D,
            ".": R
        },
        U: {
            "/": R,
            "\\": L,
            ".": U
        },
        D: {
            "/": L,
            "\\": R,
            ".": D
        }
    }
    return did[di][val]


def trace(grid, r, c, orient):
    rows = len(grid)
    cols = len(grid[0])

    ret = set()

    if orient == V:
        stack = [(r, c, U), (r, c, D)]
    else:
        stack = [(r, c, L), (r, c, R)]

    while stack:
        cr, cc, di = stack.pop()
        nr, nc = get_next(cr, cc, di)
        if not is_in(nr, nc, rows, cols) or grid[nr][nc] == "#":
            continue

        if grid[nr][nc] in ("|", "-"):
            if di in (L, R):
                ret.add((nr, nc, H))
            else:
                ret.add((nr, nc, V))
            continue

        ndi = turn(di, grid[nr][nc])
        stack.append((nr, nc, ndi))

    return ret


def show_grid(grid):
    for row in grid:
        print " ".join(row)


t = int(raw_input().strip())

for case in xrange(1, t + 1):
    rows, cols = map(int, raw_input().strip().split())
    grid = []
    for _ in xrange(rows):
        grid.append(list(raw_input().strip()))

    free = {}
    bad_beam = set()
    beam_map = dd(int)
    bad = False
    for r, c in product(xrange(rows), xrange(cols)):
        if grid[r][c] == ".":
            free[(r, c)] = set()
            for orient in (H, V):
                res = trace(grid, r, c, orient)
                free[(r, c)] |= res
        elif grid[r][c] in ('-', '|'):
            beam_map[(r, c)] = len(beam_map) + 1
            for orient in (H, V):
                res = trace(grid, r, c, orient)
                bad_beam |= res

    if bad_beam:
        beam_check = dd(int)
        for r, c, _ in bad_beam:
            beam_check[(r, c)] += 1
        if max(beam_check.values()) == 2:
            show(case, False, [])
            continue

        # Without this there can be more than
        # 2 good beams for a free field
        # although if you have more
        # than 2(x) beams for a free field
        # then there are at least x - 2
        # bad beams
        for el in bad_beam:
            for values in free.values():
                if el in values:
                    values.remove(el)

    if any(map(lambda x: len(x) not in (1, 2), free.values())):
        show(case, False, [])
        continue

    conj_normal_form = []
    for value in free.itervalues():
        vals = []
        for el in value:
            vals.append((-1 if el[2] == H else 1) * beam_map[el[:2]])
        conj_normal_form.append((vals[0], vals[1 % len(vals)]))

    for r, c, orient in bad_beam:
        sign = (-1 if orient == V else 1)  # We need to negate here
        val = sign * beam_map[(r, c)]
        conj_normal_form.append((val, val))

    possible, solution = two_sat_solve(conj_normal_form)

    if possible:
        for (r, c), val in beam_map.iteritems():
            if val in solution:
                if solution[val]:  # Means it's vertical
                    grid[r][c] = "|"
                else:
                    grid[r][c] = "-"
            else:  # Means this is isolated
                grid[r][c] = "-"
        show(case, True, grid)
    else:
        show(case, False, [])