Sorting algorithm in the element cost model Python

Question

1. Task

Using standard input I read four rows of data:

• number of elements (elements are numbered from 1 to n)
• n numbers corresponding to cost of moving each element
• n numbers corresponding to current position of each element (from 1 to n)
• n numbers corresponding to the position of each element in the desired state (from 1 to n)

I can swap every two elements I want. Cost of each operation is the sum of masses of that two elements. The program has to return minimum cost required to move elements to the desired state. The data is read using standard input (my_script.py<data.in)

2. My problem

I know the ideal theoretical solution to that problem, I created working python script but it is messy and slow. I am sure the way I convert the data to int values is stupid but I have to idea how to improve that. I am not sure if dividing list to cycles way I did is is OK, or it could be improved as well (graf is directed, but I think it doesn't matter)

method_1 and method_2 are correct mathematically, I don't won't to waste your time and try to explain why, though I am not sure if I implemented it properly.

I am new to almost every aspect of that task (data structures, standard input and so on). The code is messy.

3. Code

Input file:

10
3015 4728 4802 4361 135 4444 4313 1413 4581 546
3 10 1 8 9 4 2 7 6 5
4 9 5 3 1 6 10 7 8 2

Script I wrote:

#! /usr/bin/python3
import sys

lines = []

for line in sys.stdin:
    stripped = line.strip()
    if not stripped:
        break
    lines.append(stripped)

a, b, c, d = lines

a = a.split(" ")
b = b.split(" ")
c = c.split(" ")
d = d.split(" ")

masses = [int(i) for i in b]
unsorted = [int(i) for i in c]
sorted = [int(i) for i in d]


def dfs_test(masses, unsorted, sorted):
    segregated = []
    cycles = []
    cost = 0
    for index, element in enumerate(unsorted):
        cycle = []
        if element in segregated:
            continue
        if not sorted[index] == element:
            while element not in cycle:
                cycle.append(element)
                segregated.append(element)
                index = sorted.index(element)
                element = unsorted[index]
        else:
            cycle.append(element)
        cycles.append(cycle)
    for cycle in cycles:
        if len(cycle) == 2:
            for element in cycle:
                cost += masses[element-1]
        if len(cycle) >= 3:
            sum_mass_cycle = []
            for element in cycle:
                sum_mass_cycle.append(masses[element-1])
            method_1 = sum(sum_mass_cycle) + (len(cycle) - 2) * min(sum_mass_cycle)
            method_2 = sum(sum_mass_cycle) + min(sum_mass_cycle) + (len(cycle) + 1) * min(masses)
            cost += min(method_1, method_2)
    return cost


if __name__ == "__main__":
    print(dfs_test(masses, unsorted, sorted))

Answer 1

Style: Parsing the file

It's okay to be strict in expecting things in exactly the right format here. There should be no blank lines (or possibly, always exactly one blank line at the end) so the if check seems unneccesary. Your original is fine too, but I personally would shorten this to:

lines = [line.rstrip() for line in sys.stdin]

Drop a, b, c, and d. They are not clear names and this code is repetitive. Repetition is the enemy:

line_numbers = [[int(i) for i in line.split()] for line in line_parts]
(expected_length,), masses, unsorted, sorted = line_numbers
assert expected_length == len(masses) == len(unsorted) == len(sorted)

You can split up the definition of line_numbers into a parse_line function if that nesting is too much for you to read.

Your input is 1-indexed but python is 0-indexed. Quit using element-1 everywhere and fix it earlier on. This makes the code slightly longer but more readable and less error-prone.

sorted = [i-1 for i in sorted]
unsorted = [i-1 for i in unsorted]

sorted is the name of a built-in function, as you may be able to tell from the syntax highlighting here. It doesn't cause you any problems to use the name, but it may confuse readers. Up to you whether to keep it. A common change might be to call the variable sorted_.

All of the above should be INSIDE the main guard.

The final main guard is:

import sys

if __name__ == '__main__':
    # Parse stdin
    lines = [line.strip() for line in sys.stdin]
    line_numbers = [[int(i) for i in line.split()] for line in line_parts]
    (expected_length,), masses, unsorted, sorted = line_numbers
    assert expected_length == len(masses) == len(unsorted) == len(sorted)

    # Switch to 0-indexing
    sorted = [i-1 for i in sorted]
    unsorted = [i-1 for i in unsorted]

    print(dfs_test(masses, unsorted, sorted))

Style: dfs_test

Overall, this looks pretty clear minus the math explanations.

Overall I would say, add more comments about the intent of code (especially sections of code, ex "add the cycle containing this element"). This is not easy code.

Cycle-building:

• The cycle-building algorithm is excellent. It's clearly written and can be easily tweaked to be fast.
• I would split cycle-building into its own function that takes in unsorted and sorted and returns the cycles. This both make it clearer (provides a label for the section) and shows the reader that those two variables are the only dependencies, and aren't used elsewhere.
• segregated is not a clear name. Add a comment at the point of definitions (this is the elements already put into either cycles or the current active-build cycle.)

In the cost calculator:

• Change the if-if to if-elif, it's better style
• Make explicit the 1-length cycle case (as an empty if branch or as a comment)
• Move the definition cost = 0 to the cost section
• Consider splitting the cost calculator into its own function as well.

Algorithmic speedups

This is slower than it could be. You should know that already--don't rely on others to tell you. Measure the speed of your algorithm.

1. You can stare at it, which is actually the best method in terms of fixing or understanding a problem.
2. Or measure it experimentally, which is the best method in terms of detecting a problem. by seeing how much doubling the input size changes the runtime. If it doubles, it's O(N), if it increases it by 4X, it's O(N^2), if it does something worse think about whether you just ran out of RAM, etc. If you want to find out why something is slow experimentally, start profiling.

Now, it may be that you only are going to run this on 20-element lists, in which case you shouldn't care, and can ignore everything below. But if you want a fast runtime (if this is an algorithms problem, which it looks like it is, or if you really need to run this on large inputs in practice), you should speed it up.

• I believe this is currently an O(N^2) algorithm, and cycle-finding is the only O(N^2) step.
• sorted.index is a slow operation. List lookup does an O(N) scan, and it's being done once per element in the worst case. You can speed up cycle-building by adding a pre-processing step where you reverse the permutation to speed up this lookup. inverse_sorted = {p: i for i,p in enumerate(sorted)}.
• Change segregated to a set. if element in segregated is also O(N) time.
• These two changes should make cycle-building O(N).
• I believe cost computation should already be O(N). Speed it up by calculating sum(sum_mass_cycle) and min(sum_mass_cycle) once, instead of once for each method.

final note

all of the above code is typed directly into codereview.stackexchange with no testing, take it with a grain of salt!