4
\$\begingroup\$

I'm trying to solve the Code Jam Dice Straight problem. My Python solution seems to produce the correct output for the first few test cases but then it becomes too slow. How can I possibly optimize it?

from operator import itemgetter
from itertools import groupby    


def get_straight(chunk: list, idx: int, straight: int, dice: set) -> int:
    """Get max straight starting at index idx."""
    if idx == len(chunk) - 1:
        return straight

    item = chunk[idx]

    return max((get_straight(chunk, idx+1, straight+1, dice-{x})
                for x in item[1] if x in dice), default=straight)


def main():
    T = int(input())

    for i in range(T):
        ndice = int(input())
        integers = []

        for j in range(ndice):
            integers.extend((int(x), j) for x in input().split())

        integers.sort(key=itemgetter(0))

        grouped = [(k, [x[1] for x in g])
                   for k, g in groupby(integers, itemgetter(0))]
        chunks = []

        # find runs of consecutive numbers
        for k, g in groupby(enumerate(grouped), lambda x: x[0]-x[1][0]):
            g = list(g)
            if len(g) > 1:
                chunks.append(list(map(itemgetter(1), g)))

        chunks.sort(key=len, reverse=True)

        max_straight = 1

        for chunk in chunks:
            if len(chunk) <= max_straight:
                break

            for k in range(len(chunk)-1):
                if len(chunk) - k < max_straight:
                    break
                dice = set(range(ndice))
                straight = get_straight(chunk, k, 0, dice)
                if straight > max_straight:
                    max_straight = straight

        print("Case #{}: {}".format(i+1, max_straight))

main()
\$\endgroup\$
2
\$\begingroup\$

Your code looks good, seems to work and is somehow documented but it can still be improved in various ways (not only performances which will be only tackled at the end if at all).

Code organisation

Your call to main could be moved behind an if __name__ == "__main__": guard in case you ever want to reuse get_straight(). It may not be relevant in your case but it is a good habit to take.

Then, it may be a good idea to split your code into smaller functions, easier to understand and to maintain.

By doing so, you'd also separate the concerns my having the code computing the results you need not mixed up with the logic handling the input/output.

Finally, and maybe even more important, that will make your code testable. Indeed, before performing any changes in the algorithm, it is highly recommended to write unit-tests to prevent you from breaking everything (tests are usually not fail-proof but they help a lot).

from operator import itemgetter
from itertools import groupby    


def get_straight(chunk: list, idx: int, straight: int, dice: set) -> int:
    """Get max straight starting at index idx."""
    if idx == len(chunk) - 1:
        return straight

    item = chunk[idx]

    return max((get_straight(chunk, idx+1, straight+1, dice-{x})
                for x in item[1] if x in dice), default=straight)


def longest_straight_length(dices):
    """Return the length of the longest possible straight using `dices`."""
    integers = []
    for j, dice in enumerate(dices):
        integers.extend((val, j) for val in dice)
    integers.sort(key=itemgetter(0))

    grouped = [(k, [x[1] for x in g])
               for k, g in groupby(integers, itemgetter(0))]
    chunks = []

    # find runs of consecutive numbers
    for k, g in groupby(enumerate(grouped), lambda x: x[0]-x[1][0]):
        g = list(g)
        if len(g) > 1:
            chunks.append(list(map(itemgetter(1), g)))

    chunks.sort(key=len, reverse=True)

    max_straight = 1

    for chunk in chunks:
        if len(chunk) <= max_straight:
            break

        for k in range(len(chunk)-1):
            if len(chunk) - k < max_straight:
                break
            dice = set(range(len(dices)))
            straight = get_straight(chunk, k, 0, dice)
            if straight > max_straight:
                max_straight = straight
    return max_straight

def interactive():
    for i in range(int(input())):
        ndice = int(input())
        dices = [[int(x) for x in input().split()] for _ in range(ndice)]
        print("Case #{}: {}".format(i+1, longest_straight_length(dices)))


def unit_tests():
    test_cases = [
        ([[4, 8, 15, 16, 23, 42],
          [8, 6, 7, 5, 30, 9],
          [1, 2, 3, 4, 55, 6],
          [2, 10, 18, 36, 54, 86]],
         4),
        ([[1, 2, 3, 4, 5, 6],
          [60, 50, 40, 30, 20, 10]],
         1),
        ([[1, 2, 3, 4, 5, 6],
          [1, 2, 3, 4, 5, 6],
          [1, 4, 2, 6, 5, 3]],
         3)
    ]
    for dices, expected_result in test_cases:
        actual_result = longest_straight_length(dices)
        if actual_result != expected_result:
            print(expected_result, actual_result, dices)

if __name__ == "__main__":
    if True:
        unit_tests()
    else:
        interactive()

Rewriting the logic

The beginning of the logic creating chunks could be rewritten using a dictionnary rather than performing smart groupby and tricky index accesses. (At the end, I re-build the data structure you are using to be able to keep the code you've written). We'd get something like:

def longest_straight_length(dices):
    """Return the length of the longest possible straight using `dices`."""
    # Mapping value to list of dice indices
    number_on_dices = dict()
    for i, dice in enumerate(dices):
        for val in dice:
            number_on_dices.setdefault(val, []).append(i)

    # Consecutive chunks
    chunks = []
    for k, g in groupby(enumerate(sorted(number_on_dices.keys())), lambda x: x[0] - x[1]):
        g = list(g)
        if len(g) > 1:
            chunk = [y for (_, y) in g]
            chunks.append(chunk)
    chunks.sort(key=len, reverse=True)

    # Re-add values to be reuse existing logic for the time being
    chunks = [[(val, number_on_dices[val]) for val in chunk] for chunk in chunks]

    max_straight = 1

    for chunk in chunks:
        if len(chunk) <= max_straight:
            break

        for k in range(len(chunk)-1):
            if len(chunk) - k < max_straight:
                break
            dice = set(range(len(dices)))
            straight = get_straight(chunk, k, 0, dice)
            if straight > max_straight:
                max_straight = straight
    return max_straight

Then, it appears clearer that the first element in pairs from chunks are never used. We can simply have:

chunks = [[number_on_dices[val] for val in chunk] for chunk in chunks]

and in get_straight:

return max((get_straight(chunk, idx+1, straight+1, dice-{x})
            for x in item if x in dice), default=straight)

Ultimately, we can simplify the code building chunks:

# Consecutive chunks
chunks = []
for k, g in groupby(enumerate(sorted(number_on_dices.keys())), lambda x: x[0] - x[1]):
    g = list(g)
    if len(g) > 1:
        chunk = [number_on_dices[y] for (_, y) in g]
        chunks.append(chunk)
chunks.sort(key=len, reverse=True)

However, the expensive part is still here, in the recursive calls to get_straight. (I haven't found a trick for this yet but current version of the code is available here: http://termbin.com/yg5l ).

Bug ?

As I was trying to come up with bigger and bigger examples with more recursive calls, I've added the following test cases:

    ([[1, 2, 3, 4, 5, 6],
      [7, 2, 3, 4, 5, 6],
      [7, 8, 3, 4, 5, 6],
      [7, 8, 9, 4, 5, 6],
      [7, 8, 9,10, 5, 6],
      [7, 8, 9,10,11, 6],
      [7, 8, 9,10,11,12]],
     7),
    ([[1, 1, 1, 1, 1, 1],
      [2, 2, 2, 2, 2, 2],
      [3, 3, 3, 3, 3, 3],
      [4, 4, 4, 4, 4, 4],
      [5, 5, 5, 5, 5, 5],
      [6, 6, 6, 6, 6, 6],
      [7, 7, 7, 7, 7, 7]],
     7),

In both cases, we could generate the sequence [1, 2, 3, 4, 5, 6, 7] of length 7. However, the code returned 6. If you want to go further from that part, the code with additional tests and performance indicators is available here: http://termbin.com/c690 .

\$\endgroup\$
  • \$\begingroup\$ Thank you for your review. You're right that the recursive function is the most expensive part. I was able to speed it up a bit by turning it into a generator, as well as fix the bug you spotted and add some more optimizations. If you're interested I can update the OP with the latest version. \$\endgroup\$ – Eugene Yarmash Apr 16 '18 at 20:56
  • \$\begingroup\$ @EugeneYarmash It's best not to change the question in the question. If you want your new version of the code to be reviewed, please open a new question. Also, please make sure to take all comments into account and to include the text of the problem both in this question and in the new question. \$\endgroup\$ – Josay Apr 17 '18 at 9:21

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.