HiveBrain v1.2.0
Get Started
← Back to all entries
patternpythonMinor

HackerRank "Manasa and Stones" in Python

Submitted by: @import:stackexchange-codereview··
0
Viewed 0 times
stonesandpythonhackerrankmanasa

Problem

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.

Solution

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

Code Snippets

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

Context

StackExchange Code Review Q#61154, answer score: 7

Revisions (0)

No revisions yet.