Naming
The question starts off talking about N
and K
. Your code begins with reading in z
and d
. The question text is ready-made documentation for your code; use the same names for the variables!
n, k = map(int, input().split())
Guilty, I said N
and K
but used n
and k
. PEP 8, the Style Guide for Python Code recommends using snake_case
for variables. So we should also change S
and Flag
to s
and flag
.
But ... what is flag
? Obviously, it is a flag, but a flag for what? Perhaps changed
would be a better name; you are looping while a change was made in the previous iteration.
Similarly, S
is your list of swaps. Why not call it swaps
?
Don't relookup values
Python is an interpreted language. It doesn't do any data-flow analysis to determine if the value of an expression will be used again. It can't, since the meaning of operations can be changed on the fly.
In the for i in S
loop, how many times is i[0]
and i[1]
used? Six times each. Each time, the interpreter looks up the object stored in i
, and calls the __getitem__
method to fetch the value. Let's fix this by looking up the value only once, and storing the values in (say) p
and q
:
for p, q is swaps:
if p not in pos[permute[q - 1]]:
changed = True
pos[permute[q - 1]].append(p)
...
That is immediately clearer. No extra [0]
and [1]
to confuse the issue.
Use O(1) lookups instead of O(K) searches
You're using if thing not in container: container.append(thing)
. This is a lot like a set
. A set can only contain one instance of an item. But unlike a list, if thing not in container
is O(1) on a set, where as the with a list, the entire list has to be searched to ensure the thing is not inside the list.
Two simple changes. First, instead of pos[i] = [i,]
(the comma is actually unnecessary), initialize it with pos[i] = {i}
. Second, instead of .append()
, use .add()
to add the value to the set.
Note: sets are more complex than lists. This may actually slow things down and use more memory, if the number of things in the containers is small. The sets gain their speed advantage when they contain many elements. So profile this change, especially using larger problem sets.
Explicit is better than implicit
swaps = []
for i in range(k):
swaps.append(list(map(int, input().split())))
What happens if an input line has more than two numbers? The input will be split into whole bunch of terms, each term will be mapped (converted) to an integer, and stored as a list inside the swaps list!
Oops.
If the program crashes, it will happen well after this point in the code.
It is better to crash early. You expect exactly two values. Be explicit.
swaps = []
for i in range(k):
swaps.append(list(map(int, input().split(maxsplit=1))))
If given more than two numbers, say 12 23 34
, then split
will return "12"
and "23 34"
, and int("23 34")
will raise a ValueError
. The bug will be caught earlier.
Map the problem to Python, if possible
Python uses 0-based indexing. The problem uses 1-based indexing. You can get rid of the +1
in range
statements, and -1
in the solution code by mapping the problem into a more Pythonic version.
For example, instead of swapping numbers at positions 1
and 3
, we'll swap them at positions 0
and 2
. This just requires adjusting the swap positions as we read them in:
swaps = []
for i in range(k):
p, q = map(int, input().split())
swaps.append([p - 1, q - 1])
Note. The explicit maxsplit=1
went away because we are assigned the result of map to a tuple of two variables, which requires exactly two values. Explicit is better than implicit ... especially when the it can be expressed implicitly!
Reworked Code
n, k = map(int, input().split())
pos = {}
for i in range(n):
pos[i] = {i}
permute = list(range(n))
swaps = []
for i in range(k):
p, q = map(int, input().split())
swaps.append([p - 1, q - 1])
changed = True
while changed:
changed = False
for p, q in swaps:
if p not in pos[permute[q]]:
changed = True
pos[permute[q]].add(p)
if q not in pos[permute[p]:
changed = True
pos[permute[p]].add(q)
permute[p], permute[q] = permute[q], permute[p]
for i in pos: # Note: relies on ordered dictionary
print(len(pos[i]))
Algorithmic improvement
1 2 3 4 5
After one iteration of swaps, you end up with this:
2 4 3 1 5
We can see we have moved
- 1 -> 2, 2 -> 4, 4 -> 1. (length 3)
- 3 -> 3
- 5 -> 5
The 1 visited locations 1, 3, 2 and 4. The 2 visited locations 2, and 1. The 4 visited locations 4 and 2. On subsequent iterations, the 2 will have moved into the 1's starting position and will visit 1, 3, 2 and 4, and so on. So we have 1, 2, and 4 all visiting locations 1, 2, and 3 and 4 ... for a total of 4 locations
3 will visit 3, 1, and 2 and return to 3.
5 stays at 5.
All that is required is one pass through the swaps array to determine which positions each starting point visits, and what ending position it will have. From that, the "rings" can be determined. All starting positions in each ring will visit all the locations visited by every position in the ring.
With only one pass through the swaps array required, we don't need to save the swaps into a swaps
list; it can be processed on the fly. This helps the solution stay within the memory constraint.
n, k = map(int, input().split())
pos = [{i} for i in range(n)]
permute = list(range(n))
# Perform swaps, recording which values appear in which positions
for _ in range(k):
p, q = map(int, input().split())
p -= 1
q -= 1
pos[permute[q]].add(p)
pos[permute[p]].add(q)
permute[p], permute[q] = permute[q], permute[p]
# Implementation of final steps left to student