# Solving Code Jam's Dice Straight

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()


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 .

• 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. – Eugene Yarmash Apr 16 '18 at 20:56
• @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. – Josay Apr 17 '18 at 9:21