|
| 1 | +#ifndef ADJACENCY_LIST_GRAPH_HPP |
| 2 | +#define ADJACENCY_LIST_GRAPH_HPP |
| 3 | + |
| 4 | +#define PY_SSIZE_T_CLEAN |
| 5 | +#include <Python.h> |
| 6 | +#include <vector> |
| 7 | +#include <unordered_map> |
| 8 | +#include <string> |
| 9 | +#include "AdjacencyListGraphNode.hpp" |
| 10 | +#include "GraphEdge.hpp" |
| 11 | + |
| 12 | +extern PyTypeObject AdjacencyListGraphType; |
| 13 | + |
| 14 | +typedef struct { |
| 15 | + |
| 16 | + PyObject_HEAD |
| 17 | + std::vector<AdjacencyListGraphNode *> nodes; |
| 18 | + std::unordered_map<std::string, GraphEdge*> edges; |
| 19 | + std::unordered_map<std::string, AdjacencyListGraphNode*> node_map; |
| 20 | + |
| 21 | +} AdjacencyListGraph; |
| 22 | + |
| 23 | +static void AdjacencyListGraph_dealloc(AdjacencyListGraph* self) { |
| 24 | + for (AdjacencyListGraphNode* node : self->nodes) { |
| 25 | + Py_XDECREF(node); |
| 26 | + } |
| 27 | + self->nodes.clear(); |
| 28 | + |
| 29 | + for (auto& pair : self->edges) { |
| 30 | + Py_XDECREF(pair.second); |
| 31 | + } |
| 32 | + self->edges.clear(); |
| 33 | + |
| 34 | + self->node_map.clear(); |
| 35 | + |
| 36 | + Py_TYPE(self)->tp_free(reinterpret_cast<PyObject*>(self)); |
| 37 | +} |
| 38 | + |
| 39 | +static PyObject* AdjacencyListGraph_new(PyTypeObject* type, PyObject* args, PyObject* kwds) { |
| 40 | + AdjacencyListGraph* self = (AdjacencyListGraph*)type->tp_alloc(type, 0); |
| 41 | + if (!self) |
| 42 | + return NULL; |
| 43 | + |
| 44 | + new (&self->nodes) std::vector<AdjacencyListGraphNode*>(); |
| 45 | + new (&self->edges) std::unordered_map<std::string, GraphEdge*>(); |
| 46 | + new (&self->node_map) std::unordered_map<std::string, AdjacencyListGraphNode*>(); |
| 47 | + |
| 48 | + Py_ssize_t num_args = PyTuple_Size(args); |
| 49 | + for (Py_ssize_t i = 0; i < num_args; ++i) { |
| 50 | + PyObject* node_obj = PyTuple_GetItem(args, i); |
| 51 | + if (!PyObject_IsInstance(node_obj, (PyObject*)&AdjacencyListGraphNodeType)) { |
| 52 | + PyErr_SetString(PyExc_TypeError, "All arguments must be AdjacencyListGraphNode instances"); |
| 53 | + |
| 54 | + self->nodes.~vector(); |
| 55 | + self->edges.~unordered_map(); |
| 56 | + self->node_map.~unordered_map(); |
| 57 | + Py_TYPE(self)->tp_free((PyObject*)self); |
| 58 | + return NULL; |
| 59 | + } |
| 60 | + |
| 61 | + AdjacencyListGraphNode* node = (AdjacencyListGraphNode*)node_obj; |
| 62 | + |
| 63 | + if (self->node_map.find(node->name) != self->node_map.end()) { |
| 64 | + PyErr_Format(PyExc_ValueError, "Duplicate node with name '%s'", node->name.c_str()); |
| 65 | + |
| 66 | + self->nodes.~vector(); |
| 67 | + self->edges.~unordered_map(); |
| 68 | + self->node_map.~unordered_map(); |
| 69 | + Py_TYPE(self)->tp_free((PyObject*)self); |
| 70 | + return NULL; |
| 71 | + } |
| 72 | + |
| 73 | + Py_INCREF(node); |
| 74 | + self->nodes.push_back(node); |
| 75 | + self->node_map[node->name] = node; |
| 76 | + } |
| 77 | + return (PyObject*)self; |
| 78 | +} |
| 79 | + |
| 80 | +static std::string make_edge_key(const std::string& source, const std::string& target) { |
| 81 | + return source + "_" + target; |
| 82 | +} |
| 83 | + |
| 84 | +static PyObject* AdjacencyListGraph_add_vertex(AdjacencyListGraph* self, PyObject* args) { |
| 85 | + PyObject* node_obj; |
| 86 | + |
| 87 | + if (!PyArg_ParseTuple(args, "O", &node_obj)) { |
| 88 | + PyErr_SetString(PyExc_ValueError, "Expected a single AdjacencyListGraphNode object"); |
| 89 | + return NULL; |
| 90 | + } |
| 91 | + |
| 92 | + if (!PyObject_IsInstance(node_obj, (PyObject*)&AdjacencyListGraphNodeType)) { |
| 93 | + PyErr_SetString(PyExc_TypeError, "Object is not an AdjacencyListGraphNode"); |
| 94 | + return NULL; |
| 95 | + } |
| 96 | + |
| 97 | + AdjacencyListGraphNode* node = (AdjacencyListGraphNode*)node_obj; |
| 98 | + |
| 99 | + if (self->node_map.find(node->name) != self->node_map.end()) { |
| 100 | + PyErr_SetString(PyExc_ValueError, "Node with this name already exists"); |
| 101 | + return NULL; |
| 102 | + } |
| 103 | + |
| 104 | + Py_INCREF(node); |
| 105 | + self->nodes.push_back(node); |
| 106 | + self->node_map[node->name] = node; |
| 107 | + |
| 108 | + Py_RETURN_NONE; |
| 109 | +} |
| 110 | + |
| 111 | + |
| 112 | +static PyObject* AdjacencyListGraph_is_adjacent(AdjacencyListGraph* self, PyObject* args) { |
| 113 | + const char* node1_name_c; |
| 114 | + const char* node2_name_c; |
| 115 | + if (!PyArg_ParseTuple(args, "ss", &node1_name_c, &node2_name_c)) |
| 116 | + return NULL; |
| 117 | + |
| 118 | + std::string node1_name(node1_name_c); |
| 119 | + std::string node2_name(node2_name_c); |
| 120 | + |
| 121 | + auto it1 = self->node_map.find(node1_name); |
| 122 | + if (it1 == self->node_map.end()) { |
| 123 | + PyErr_SetString(PyExc_KeyError, "node1 not found"); |
| 124 | + return NULL; |
| 125 | + } |
| 126 | + AdjacencyListGraphNode* node1 = it1->second; |
| 127 | + |
| 128 | + if (node1->adjacent.find(node2_name) != node1->adjacent.end()) { |
| 129 | + Py_RETURN_TRUE; |
| 130 | + } else { |
| 131 | + Py_RETURN_FALSE; |
| 132 | + } |
| 133 | +} |
| 134 | + |
| 135 | +static PyObject* AdjacencyListGraph_num_vertices(AdjacencyListGraph* self, PyObject* Py_UNUSED(ignored)) { |
| 136 | + return PyLong_FromSize_t(self->nodes.size()); |
| 137 | +} |
| 138 | + |
| 139 | +static PyObject* AdjacencyListGraph_num_edges(AdjacencyListGraph* self, PyObject* Py_UNUSED(ignored)) { |
| 140 | + return PyLong_FromSize_t(self->edges.size()); |
| 141 | +} |
| 142 | + |
| 143 | +static PyObject* AdjacencyListGraph_neighbors(AdjacencyListGraph* self, PyObject* args) { |
| 144 | + const char* node_name_c; |
| 145 | + if (!PyArg_ParseTuple(args, "s", &node_name_c)) |
| 146 | + return NULL; |
| 147 | + |
| 148 | + std::string node_name(node_name_c); |
| 149 | + auto it = self->node_map.find(node_name); |
| 150 | + if (it == self->node_map.end()) { |
| 151 | + PyErr_SetString(PyExc_KeyError, "Node not found"); |
| 152 | + return NULL; |
| 153 | + } |
| 154 | + AdjacencyListGraphNode* node = it->second; |
| 155 | + |
| 156 | + PyObject* neighbors_list = PyList_New(0); |
| 157 | + if (!neighbors_list) return NULL; |
| 158 | + |
| 159 | + for (const auto& adj_pair : node->adjacent) { |
| 160 | + Py_INCREF(adj_pair.second); |
| 161 | + PyList_Append(neighbors_list, (PyObject*)adj_pair.second); |
| 162 | + } |
| 163 | + |
| 164 | + return neighbors_list; |
| 165 | +} |
| 166 | + |
| 167 | +static PyObject* AdjacencyListGraph_remove_vertex(AdjacencyListGraph* self, PyObject* args) { |
| 168 | + const char* name_c; |
| 169 | + if (!PyArg_ParseTuple(args, "s", &name_c)) |
| 170 | + return NULL; |
| 171 | + |
| 172 | + std::string name(name_c); |
| 173 | + auto it = self->node_map.find(name); |
| 174 | + if (it == self->node_map.end()) { |
| 175 | + PyErr_SetString(PyExc_KeyError, "Node not found"); |
| 176 | + return NULL; |
| 177 | + } |
| 178 | + |
| 179 | + AdjacencyListGraphNode* node_to_remove = it->second; |
| 180 | + |
| 181 | + auto& nodes_vec = self->nodes; |
| 182 | + nodes_vec.erase(std::remove(nodes_vec.begin(), nodes_vec.end(), node_to_remove), nodes_vec.end()); |
| 183 | + |
| 184 | + self->node_map.erase(it); |
| 185 | + |
| 186 | + Py_XDECREF(node_to_remove); |
| 187 | + |
| 188 | + for (auto& node_pair : self->node_map) { |
| 189 | + AdjacencyListGraphNode* node = node_pair.second; |
| 190 | + auto adj_it = node->adjacent.find(name); |
| 191 | + if (adj_it != node->adjacent.end()) { |
| 192 | + Py_XDECREF(adj_it->second); |
| 193 | + node->adjacent.erase(adj_it); |
| 194 | + } |
| 195 | + } |
| 196 | + |
| 197 | + for (auto it = self->edges.begin(); it != self->edges.end(); ) { |
| 198 | + const std::string& key = it->first; |
| 199 | + |
| 200 | + bool involves_node = false; |
| 201 | + size_t underscore = key.find('_'); |
| 202 | + if (underscore != std::string::npos) { |
| 203 | + std::string source = key.substr(0, underscore); |
| 204 | + std::string target = key.substr(underscore + 1); |
| 205 | + if (source == name || target == name) |
| 206 | + involves_node = true; |
| 207 | + } |
| 208 | + |
| 209 | + if (involves_node) { |
| 210 | + Py_XDECREF(it->second); |
| 211 | + it = self->edges.erase(it); |
| 212 | + } else { |
| 213 | + ++it; |
| 214 | + } |
| 215 | + } |
| 216 | + |
| 217 | + Py_RETURN_NONE; |
| 218 | +} |
| 219 | + |
| 220 | + |
| 221 | +static PyObject* AdjacencyListGraph_get_edge(AdjacencyListGraph* self, PyObject* args) { |
| 222 | + const char* source_c; |
| 223 | + const char* target_c; |
| 224 | + if (!PyArg_ParseTuple(args, "ss", &source_c, &target_c)) |
| 225 | + return NULL; |
| 226 | + |
| 227 | + std::string key = make_edge_key(source_c, target_c); |
| 228 | + auto it = self->edges.find(key); |
| 229 | + if (it != self->edges.end()) { |
| 230 | + Py_INCREF(it->second); |
| 231 | + return (PyObject*)it->second; |
| 232 | + } |
| 233 | + |
| 234 | + Py_RETURN_NONE; |
| 235 | +} |
| 236 | + |
| 237 | +static PyObject* AdjacencyListGraph_remove_edge(AdjacencyListGraph* self, PyObject* args) { |
| 238 | + const char* source_c; |
| 239 | + const char* target_c; |
| 240 | + if (!PyArg_ParseTuple(args, "ss", &source_c, &target_c)) |
| 241 | + return NULL; |
| 242 | + |
| 243 | + std::string source(source_c); |
| 244 | + std::string target(target_c); |
| 245 | + |
| 246 | + auto it_source = self->node_map.find(source); |
| 247 | + auto it_target = self->node_map.find(target); |
| 248 | + if (it_source == self->node_map.end() || it_target == self->node_map.end()) { |
| 249 | + PyErr_SetString(PyExc_KeyError, "Source or target node not found"); |
| 250 | + return NULL; |
| 251 | + } |
| 252 | + |
| 253 | + AdjacencyListGraphNode* source_node = it_source->second; |
| 254 | + |
| 255 | + auto adj_it = source_node->adjacent.find(target); |
| 256 | + if (adj_it != source_node->adjacent.end()) { |
| 257 | + Py_XDECREF(adj_it->second); |
| 258 | + source_node->adjacent.erase(adj_it); |
| 259 | + } |
| 260 | + |
| 261 | + std::string key = make_edge_key(source, target); |
| 262 | + auto edge_it = self->edges.find(key); |
| 263 | + if (edge_it != self->edges.end()) { |
| 264 | + Py_XDECREF(edge_it->second); |
| 265 | + self->edges.erase(edge_it); |
| 266 | + } |
| 267 | + |
| 268 | + Py_RETURN_NONE; |
| 269 | +} |
| 270 | + |
| 271 | +static PyObject* AdjacencyListGraph_add_edge(AdjacencyListGraph* self, PyObject* args) { |
| 272 | + const char* source_cstr; |
| 273 | + const char* target_cstr; |
| 274 | + PyObject* value = Py_None; |
| 275 | + |
| 276 | + if (!PyArg_ParseTuple(args, "ss|O", &source_cstr, &target_cstr, &value)) { |
| 277 | + PyErr_SetString(PyExc_ValueError, "Expected source (str), target (str), and optional weight"); |
| 278 | + return NULL; |
| 279 | + } |
| 280 | + |
| 281 | + std::string source(source_cstr); |
| 282 | + std::string target(target_cstr); |
| 283 | + |
| 284 | + if (self->node_map.find(source) == self->node_map.end()) { |
| 285 | + PyErr_SetString(PyExc_ValueError, "Source node does not exist"); |
| 286 | + return NULL; |
| 287 | + } |
| 288 | + if (self->node_map.find(target) == self->node_map.end()) { |
| 289 | + PyErr_SetString(PyExc_ValueError, "Target node does not exist"); |
| 290 | + return NULL; |
| 291 | + } |
| 292 | + |
| 293 | + AdjacencyListGraphNode* source_node = self->node_map[source]; |
| 294 | + AdjacencyListGraphNode* target_node = self->node_map[target]; |
| 295 | + |
| 296 | + PyObject* edge_args = PyTuple_Pack(3, (PyObject*)source_node, (PyObject*)target_node, value); |
| 297 | + if (!edge_args) return NULL; |
| 298 | + |
| 299 | + PyObject* edge_obj = PyObject_CallObject((PyObject*)&GraphEdgeType, edge_args); |
| 300 | + Py_DECREF(edge_args); |
| 301 | + |
| 302 | + if (!edge_obj) |
| 303 | + return NULL; |
| 304 | + |
| 305 | + std::string key = make_edge_key(source, target); |
| 306 | + |
| 307 | + auto it = self->edges.find(key); |
| 308 | + if (it != self->edges.end()) { |
| 309 | + Py_XDECREF(it->second); |
| 310 | + it->second = (GraphEdge*)edge_obj; |
| 311 | + } else { |
| 312 | + self->edges[key] = (GraphEdge*)edge_obj; |
| 313 | + } |
| 314 | + |
| 315 | + auto adj_it = source_node->adjacent.find(target); |
| 316 | + if (adj_it == source_node->adjacent.end()) { |
| 317 | + Py_INCREF(target_node); |
| 318 | + source_node->adjacent[target] = (PyObject*)target_node; |
| 319 | + } |
| 320 | + |
| 321 | + Py_RETURN_NONE; |
| 322 | +} |
| 323 | + |
| 324 | + |
| 325 | + |
| 326 | + |
| 327 | +static PyMethodDef AdjacencyListGraph_methods[] = { |
| 328 | + {"add_vertex", (PyCFunction)AdjacencyListGraph_add_vertex, METH_VARARGS, "Add a vertex to the graph"}, |
| 329 | + {"add_edge", (PyCFunction)AdjacencyListGraph_add_edge, METH_VARARGS, "Add an edge to the graph"}, |
| 330 | + {"is_adjacent", (PyCFunction)AdjacencyListGraph_is_adjacent, METH_VARARGS, "Check adjacency between two nodes"}, |
| 331 | + {"num_vertices", (PyCFunction)AdjacencyListGraph_num_vertices, METH_NOARGS, "Number of vertices"}, |
| 332 | + {"num_edges", (PyCFunction)AdjacencyListGraph_num_edges, METH_NOARGS, "Number of edges"}, |
| 333 | + {"neighbors", (PyCFunction)AdjacencyListGraph_neighbors, METH_VARARGS, "Get neighbors of a node"}, |
| 334 | + {"remove_vertex", (PyCFunction)AdjacencyListGraph_remove_vertex, METH_VARARGS, "Remove a vertex"}, |
| 335 | + {"get_edge", (PyCFunction)AdjacencyListGraph_get_edge, METH_VARARGS, "Get edge between source and target"}, |
| 336 | + {"remove_edge", (PyCFunction)AdjacencyListGraph_remove_edge, METH_VARARGS, "Remove edge between source and target"}, |
| 337 | + {NULL} |
| 338 | +}; |
| 339 | + |
| 340 | + |
| 341 | +PyTypeObject AdjacencyListGraphType = { |
| 342 | + PyVarObject_HEAD_INIT(NULL, 0) |
| 343 | + .tp_init= 0, |
| 344 | + .tp_name = "_graph.AdjacencyListGraph", |
| 345 | + .tp_basicsize = sizeof(AdjacencyListGraph), |
| 346 | + .tp_itemsize = 0, |
| 347 | + .tp_dealloc = (destructor)AdjacencyListGraph_dealloc, |
| 348 | + .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, |
| 349 | + .tp_doc = "Adjacency List Graph data structure", |
| 350 | + .tp_methods = AdjacencyListGraph_methods, |
| 351 | + .tp_new = AdjacencyListGraph_new, |
| 352 | +}; |
| 353 | + |
| 354 | +#endif |
| 355 | + |
0 commit comments