[PYTHON] Basics of computational complexity improvement learned from ABC163C

Introduction

When a friend who had almost never touched programming told me "Tell me Python!" And held a study session, ABC163C --management Was screaming, "It's a good question!"

In fact, I was interested in the shift from the phase of "being able to write code" to the phase of "being able to improve the amount of calculation" in one problem, so I will try to describe the interaction.

Who is the target of this article

Someone who remembers grammar such as for, if, and list to some extent and is confident in logic but somehow becomes TLE.

Problem summary

Print the number of direct reports for employee numbers $ 1, 2, \ cdots, and N $. $ A_2, \ cdots, A_N $ is given, where the boss of the $ i $ th employee is $ A_i $. However, the employee number of the boss is younger than the employee number of a certain employee.

Trajectory of computational complexity improvement

It was TLE twice and AC the third time.

1. First of all, be honest

policy

Search the contents of boss list A in order from employee number 1 to N and print the number of hits. The count () method can be used to output the specified number of elements for the list, so use that.

abc163_c.py

N = int(input())
A = list(map(int, input().split()))

for i in range(1, N + 1):
    count = A.count(i)
    print(count)

Submit!

TLE! → Why not?

This is because the constraint is $ 2 \ le N \ le 2 \ times 10 ^ 5 $, but the amount of calculation is $ O (N ^ 2) \ approx 10 ^ {10} $.

The method count () counts the elements of the list, comparing them end-to-end.

Assuming that the number of elements in the list is N, [calculate N times](#Supplementary listcount implementation) every time count () is called. It's $ O (N ^ 2) $ because it's written in a for statement that loops N times.

In addition, in the programming contest challenge book (commonly known as ant book), if the time limit is 1 second in C ++ which is the compilation language

$ 10 ^ 6 $ In time with a margin $ 10 ^ 7 $ Probably in time $ 10 ^ 8 $ Strict unless it is a very simple process

There is. The amount of calculation $ 10 ^ {10} $ in Python is too late.

2. Try to stop in the middle

policy

When the sum of the number of subordinates obtained by count () reaches $ n -1 $, there is no need to rely on this method anymore. Therefore, define the variable total that holds the sum with an initial value of 0, and add the results each time you callcount ().

However, since it is not known how many times the surplus employee 0 should be output, the employee number at the time of termination is saved in a variable called t.

abc163_c.py

N = int(input())
A = list(map(int, input().split()))

total = 0
t = 0
for i in range(1, N):
    count = A.count(i)
    print(count)
    total += count
    if total == N - 1:
        t = i
        break
for _ in range(t, N):
    print(0)

Submit!

The second TLE! → Is this still useless?

This is because ** worst computational complexity ** becomes $ O (N ^ 2) $. For example

7
1 2 3 4 5 6

In that case, you can finally break the for loop that uses count () when you see employee number 1 and up to 6. In this way, assuming the worst case where the calculation does not end without looking at employee $ 1, \ cdots, N -1 $, it is $ O (N ^ 2) $ in the end.

So, it seems very useless to calculate the number of subordinates for each employee every time.

3. Save the result in another array

policy

Until now, I used to calculate the number of direct reports for each employee, but I wondered if I could somehow make a list containing the number of subordinates of each employee with $ O (N) $. I will try. Finally, the contents of the list should be output one by one.

To do this, you need to keep track of how many times you see each employee number as you walk through the list of your direct supervisors $ A $. You need something that can hold data for $ N $ as a recording destination [^ 1], but it has a length of $ N $ such that the index corresponds to each employee and the element is the number of subordinates. It looks good to keep it as a list.

In the figure, it looks like this. The upper card is List A, and the lower container is "Recording destination".

[^ 1]: Due to the limitation of the problem, the Nth employee cannot be the boss, so the value for the Nth employee in this array is always 0, but considering the time and effort when outputting, N people Minutes.

abc163_c.py

N = int(input())
A = list(map(int, input().split()))

B = [0] * N
for i in A:
    B[i - 1] += 1
for i in B:
    print(i)

• Congratulations AC![^ 2] *

[^ 2]: Since we are not accessing the element of A, we can keep it in the state of generator at the time of standard input ʻA = map (int, input (). Split ()))`.

at the end

I think it's really important to estimate that this amount of calculation is not good and think of a different policy.

Supplement / About the implementation of list.count

Excerpt from cpython / Objects what the implementation of list.count is.

listobject.c

/*[clinic input]
list.count

     value: object
     /

Return number of occurrences of value.
[clinic start generated code]*/

static PyObject *
list_count(PyListObject *self, PyObject *value)
/*[clinic end generated code: output=b1f5d284205ae714 input=3bdc3a5e6f749565]*/
{
    Py_ssize_t count = 0;
    Py_ssize_t i;

    for (i = 0; i < Py_SIZE(self); i++) {
        PyObject *obj = self->ob_item[i];
        if (obj == value) {
           count++;
           continue;
        }
        Py_INCREF(obj);
        int cmp = PyObject_RichCompareBool(obj, value, Py_EQ);
        Py_DECREF(obj);
        if (cmp > 0)
            count++;
        else if (cmp < 0)
            return NULL;
    }
    return PyLong_FromSsize_t(count);
}

I don't usually write C language, so I won't discuss it rigorously,

--Look at the list end-to-end with a for statement: for (i = 0; i <Py_SIZE (self); i ++) --Check the match with the if statement: ʻif (obj == value) --If they match, advance the counter by 1:count ++;`

There are. So it seems that list.count can be said to be O (N).