summaryrefslogtreecommitdiff
path: root/vendor/github.com/DavidBelicza/TextRank/v2/rank/relation.go
blob: cb8b97ed7efb9154868ab32715cfb15dfed0596a (plain)
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
package rank

// Relation struct contains the phrase data.
//
// Max is the occurrence of the most used phrase.
//
// Min is the occurrence of the less used phrase. It is always greater then 0.
//
// Node is contains the Scores. Firs ID is the word 1, second ID is the word 2,
// and the value is the Score what contains the data about their relation.
type Relation struct {
	Max  float32
	Min  float32
	Node map[int]map[int]Score
}

// Score struct contains data about a relation of two words.
//
// Qty is the occurrence of the phrase.
//
// Weight is the weight of the phrase between 0.00 and 1.00.
//
// SentenceIDs contains all IDs of sentences what contain the phrase.
type Score struct {
	Qty         int
	Weight      float32
	SentenceIDs []int
}

// AddRelation method adds a new relation to Relation object.
func (relation *Relation) AddRelation(wordID int, relatedWordID int, sentenceID int) {
	if relatedWordID == -1 {
		return
	}

	if relation.updateRelation(relatedWordID, wordID, true, sentenceID) {
		return
	}

	if relation.extendRelation(wordID, relatedWordID, true, sentenceID) {
		return
	}

	relation.createRelation(wordID, relatedWordID, sentenceID)
}

func (relation *Relation) updateRelation(x int, y int, r bool, sentenceID int) bool {
	if _, ok := relation.Node[x][y]; ok {
		count := relation.Node[x][y].Qty + 1
		weight := relation.Node[x][y].Weight
		sentenceIDs := append(relation.Node[x][y].SentenceIDs, sentenceID)
		relation.Node[x][y] = Score{count, weight, sentenceIDs}

		return true
	} else if r {
		return relation.updateRelation(y, x, false, sentenceID)
	}

	return false
}

func (relation *Relation) extendRelation(x int, y int, r bool, sentenceID int) bool {
	if _, ok := relation.Node[x]; ok {
		relation.Node[x][y] = Score{1, 0, []int{sentenceID}}

		return true
	} else if r {
		return relation.extendRelation(y, x, false, sentenceID)
	}

	return false
}

func (relation *Relation) createRelation(x int, y int, sentenceID int) {
	relation.Node[x] = map[int]Score{}
	relation.Node[x][y] = Score{1, 0, []int{sentenceID}}
}