Как переписать этот метод с C++ на Python?

Есть метод на C++

vector<unordered_set<int>> sortGraph(vector<vector<int>> &g, vector<int> &p) {
    vector<unordered_set<int>> answer;
    for (int i = 0; i < g.size(); i++)
        answer.push_back(unordered_set<int>());
    for (int i = 0; i < g.size(); i++)
        answer[p[i]].insert(i);
    return answer;
}

Правильно ли я переписал его на Python?

def sort_graph(g, p):
    answer = []
    for i in range(len(g)):
        answer.append(set())
    for i in range(len(g)):
        answer[p[i]].add(i)
    return answer 

Еще почему-то во втором цикле в add i подчеркивается и IDE пишет Expected type 'set' (matched generic type '_T'), got 'int' instead