graph_algorithms/graph.py at master · dkositsyn/graph_algorithms · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

import copy
import itertools
import collections


class ArcMarks(collections.Mapping):
    def __init__(self):
        self._marks = {}

    def copy(self):
        return copy.deepcopy(self)

    def get_mark(self, vertex_from, vertex_to):
        return self._marks.get((vertex_from, vertex_to), None)

    def set_mark(self, vertex_from, vertex_to, value):
        self._marks[(vertex_from, vertex_to)] = value

    def del_mark(self, vertex_from, vertex_to):
        if (vertex_from, vertex_to) in self._marks:
            del self._marks[(vertex_from, vertex_to)]

    def reset_marks(self, default_value=None):
        for key in self._marks.iterkeys():
            self._marks[key] = default_value

    def __getitem__(self, key):
        vertex_from, vertex_to = key
        return self.get_mark(vertex_from, vertex_to)

    def __setitem__(self, key, value):
        vertex_from, vertex_to = key
        self.set_mark(vertex_from, vertex_to, value)

    def __delitem__(self, key):
        vertex_from, vertex_to = key
        self.del_mark(vertex_from, vertex_to)

    def __iter__(self):
        return iter(self._marks)

    def __len__(self):
        return len(self._marks)


class Graph(object):
    def __init__(self, vertex_count):
        self._arcs = [set() for _ in xrange(vertex_count)]  # doesn't support multiple arcs
        self._reversed_arcs = [set() for _ in xrange(vertex_count)]
        self._marks = ArcMarks()

    @property
    def arc_count(self):
        return sum(map(len, self._arcs))

    def __iter__(self):
        for adjacent_arcs in self._arcs:
            yield adjacent_arcs

    @property
    def vertex_count(self):
        return len(self._arcs)

    def __len__(self):
        return self.vertex_count

    def get_forward(self, vertex_idx):
        return self._arcs[vertex_idx]

    def get_backward(self, vertex_idx):
        return self._reversed_arcs[vertex_idx]

    def __getitem__(self, vertex_idx):
        return self.get_forward(vertex_idx)

    def has(self, vertex_from, vertex_to):
        return vertex_to in self[vertex_from]

    def get_mark_collection(self):
        """:return ArcMarks"""
        return self._marks.copy()

    def get_mark(self, vertex_from, vertex_to):
        return self._marks.get_mark(vertex_from, vertex_to)

    def set_mark(self, vertex_from, vertex_to, value):
        self._marks.set_mark(vertex_from, vertex_to, value)

    def add(self, vertex_from, vertex_to, value=None):
        if vertex_from != vertex_to:  # doesn't support loops
            self._arcs[vertex_from].add(vertex_to)
            self._reversed_arcs[vertex_to].add(vertex_from)
            if value is not None:
                self._marks.set_mark(vertex_from, vertex_to, value)

    def remove(self, vertex_from, vertex_to):
        self._arcs[vertex_from].remove(vertex_to)
        self._reversed_arcs[vertex_to].remove(vertex_from)
        self._marks.del_mark(vertex_from, vertex_to)

    def copy(self):
        """:return Graph"""
        return copy.deepcopy(self)


class UndirectedGraph(Graph):
    @property
    def edge_count(self):
        return self.arc_count

    def __iter__(self):
        for (adj_arcs, adj_backward_arcs) in itertools.izip(self._arcs, self._reversed_arcs):
            yield itertools.chain(adj_arcs, adj_backward_arcs)

    def get_adjacent(self, vertex_idx):
        return itertools.chain(self.get_forward(vertex_idx), self.get_backward(vertex_idx))

    def __getitem__(self, vertex_idx):
        return self.get_adjacent(vertex_idx)

    def get_mark(self, vertex_from, vertex_to):
        if vertex_to not in self.get_forward(vertex_from):
            vertex_from, vertex_to = vertex_to, vertex_from
        mark = super(UndirectedGraph, self).get_mark(vertex_from, vertex_to)
        return mark

    def set_mark(self, vertex_from, vertex_to, value):
        if vertex_to not in self.get_forward(vertex_from):
            vertex_from, vertex_to = vertex_to, vertex_from
        super(UndirectedGraph, self).set_mark(vertex_from, vertex_to, value)