HackerRank "Manasa and Stones" in Python

Question

Problem Statement:

Manasa is out on a hike with friends. She finds a trail of stones with numbers on them. She starts following the trail and notices that two consecutive stones have a difference of either a or b. Legend has it that there is a treasure trove at the end of the trail and if Manasa can guess the value of the last stone, the treasure would be hers. Given that the number on the first stone was 0, find all the possible values for the number on the last stone.

Note : The numbers on the stones are in increasing order

Input Format:

The first line contains an integer T i.e. the number of Test cases. T testcases follow. Each testcase has 3 lines. The first line contains n ( the number of stones ) The second line contains a. The third line contains b.

Output Format:

Space separated list of numbers which are the possible values of the last stone in increasing order.

Constraints:

1 ≤ T ≤ 10

1 ≤ n, a, b ≤ 103

Algorithm:

I order a and b such that a > b. The maximal value of the last stone is then found from the sequence:

1, a, 2 * a, 3 * a, ..., (n-2) * a, (n-1) * a

Which is arithmetic sequence with difference a. A smaller value will be produced by replacing the last difference of 'a' with 'b', which results in (n-2) * a + 1 * b. Doing this repeatedly yields the list of possible stone values:

(n-1) * a + 0 * b, (n-2) * a + 1 * b, ..., 1 * a + (n-2) * b, 0 * a + (n-1) * b

Since a, b can be equal, I wrap this list into a set to produce my code below.

Code:

T = int(input())

for _ in range(T):
    n = int(input())
    a = int(input())
    b = int(input())
    a, b = sorted((a, b))
    values = sorted(set([(n - i - 1) * a + i * b for i in range(n)]))
    print(*values) 

I would like any suggestions to help improve this code.

Answer 1

Since a, b can be equal, I wrap this list into a set to produce my code below.

Why not just handle this case separately? It's then simple to prove that if a != b, we don't have repeated values, and we can just generate them in increasing order.

def last_stone(a, b, n):
    if a == b:
        return [a * (n - 1)]
    if a < b:
        return last_stone(b, a, n)
    return [i * a + (n - i - 1) * b for i in range(n)]

Answer 2

Write the algorithm as a generator so that it outputs the numbers in monotonically increasing order, as required by the answer. Then you can drop successive same counts by using itertools.groupby.

itertools.groupby returns a generator where successive values are pairs (k, g) where k is the key for the group and g is an iterator for the group.

from itertools import groupby
T = int(input())

for _ in range(T):
    n = int(input())
    ab = int(input()), int(input())
    a, b = min(ab), max(ab)

    stonegen = ((n - i - 1) * a + i * b for i in range(n))
    print(*(k for k, g in groupby(stonegen)))