|
| 1 | +package data_structures |
| 2 | + |
| 3 | +import ( |
| 4 | + "sort" |
| 5 | + |
| 6 | + testerutils_random "github.com/codecrafters-io/tester-utils/random" |
| 7 | +) |
| 8 | + |
| 9 | +type SortedSetMember struct { |
| 10 | + Name string |
| 11 | + Score float64 |
| 12 | +} |
| 13 | + |
| 14 | +// SortedSet is a data structure that maintains its elements sorted by score. |
| 15 | +// If multiple elements have the same score, they are ordered lexicographically. |
| 16 | +type SortedSet struct { |
| 17 | + members []SortedSetMember |
| 18 | +} |
| 19 | + |
| 20 | +func NewSortedSet() *SortedSet { |
| 21 | + return &SortedSet{} |
| 22 | +} |
| 23 | + |
| 24 | +func (ss *SortedSet) AddMember(m SortedSetMember) *SortedSet { |
| 25 | + ss.members = append(ss.members, m) |
| 26 | + ss.sort() |
| 27 | + return ss |
| 28 | +} |
| 29 | + |
| 30 | +func (ss *SortedSet) RemoveMember(name string) *SortedSet { |
| 31 | + for i, m := range ss.members { |
| 32 | + if m.Name == name { |
| 33 | + ss.members = append(ss.members[:i], ss.members[i+1:]...) |
| 34 | + return ss |
| 35 | + } |
| 36 | + } |
| 37 | + return ss |
| 38 | +} |
| 39 | + |
| 40 | +func (ss *SortedSet) Size() int { |
| 41 | + return len(ss.members) |
| 42 | +} |
| 43 | + |
| 44 | +// GetMembers returns the a copy of all the members |
| 45 | +func (ss *SortedSet) GetMembers() []SortedSetMember { |
| 46 | + members := make([]SortedSetMember, len(ss.members)) |
| 47 | + copy(members, ss.members) |
| 48 | + return members |
| 49 | +} |
| 50 | + |
| 51 | +// GetMemberNames returns a slice containing all member names |
| 52 | +func (ss *SortedSet) GetMemberNames() []string { |
| 53 | + memberNames := make([]string, len(ss.members)) |
| 54 | + for i, m := range ss.members { |
| 55 | + memberNames[i] = m.Name |
| 56 | + } |
| 57 | + return memberNames |
| 58 | +} |
| 59 | + |
| 60 | +type SortedSetMemberGenerationOption struct { |
| 61 | + Count int // Total number of members to generate |
| 62 | + SameScoreCount int // Number of members with same score (for testing lexicographic sorting) |
| 63 | +} |
| 64 | + |
| 65 | +func GenerateSortedSetWithRandomMembers(option SortedSetMemberGenerationOption) *SortedSet { |
| 66 | + count := option.Count |
| 67 | + sameScoreCount := min(option.SameScoreCount, count) |
| 68 | + differentScoresCount := count - sameScoreCount |
| 69 | + |
| 70 | + ss := NewSortedSet() |
| 71 | + |
| 72 | + memberNames := testerutils_random.RandomWords(count) |
| 73 | + |
| 74 | + // generate members with different scores |
| 75 | + for i := range differentScoresCount { |
| 76 | + score := GetRandomSortedSetScore() |
| 77 | + ss.AddMember(SortedSetMember{ |
| 78 | + Name: memberNames[i], |
| 79 | + Score: score, |
| 80 | + }) |
| 81 | + } |
| 82 | + |
| 83 | + // generate members with same score |
| 84 | + baseScore := GetRandomSortedSetScore() |
| 85 | + for i := range sameScoreCount { |
| 86 | + ss.AddMember(SortedSetMember{ |
| 87 | + Name: memberNames[differentScoresCount+i], |
| 88 | + Score: baseScore, |
| 89 | + }) |
| 90 | + } |
| 91 | + |
| 92 | + return ss |
| 93 | +} |
| 94 | + |
| 95 | +// GetRandomSortedSetScore returns a random value of score for a sorted set |
| 96 | +func GetRandomSortedSetScore() float64 { |
| 97 | + return testerutils_random.RandomFloat64(1, 100) |
| 98 | +} |
| 99 | + |
| 100 | +// sort orders members by ascending value of score |
| 101 | +// if scores are same, the members are sorted lexicographically |
| 102 | +func (ss *SortedSet) sort() { |
| 103 | + sort.Slice(ss.members, func(i, j int) bool { |
| 104 | + if ss.members[i].Score != ss.members[j].Score { |
| 105 | + return ss.members[i].Score < ss.members[j].Score |
| 106 | + } |
| 107 | + return ss.members[i].Name < ss.members[j].Name |
| 108 | + }) |
| 109 | +} |
0 commit comments